当前位置：首页 > news >正文

【保姆级视频教程（二）】YOLOv12训练数据集构建：标签格式转换-划分-YAML 配置避坑指南 | 小白也能轻松玩转目标检测！

news 来源：原创 2025/5/23 10:40:52

【2025全站首发】YOLOv12训练数据集构建：标签格式转换-划分-YAML 配置避坑指南 | 小白也能轻松玩转目标检测！

文章目录

1. 数据集准备
- 1.1 标签格式转换
- 1.2 数据集划分
- 1.3 yaml配置文件创建
2. 训练验证

1. 数据集准备

示例数据集下载链接：PKU-Market-PCB数据集

1.1 标签格式转换

cursor prompt

请撰写一个py脚本。将@Annotations文件夹下的所有类别的xml格式的标注转换成为yolo格式。转换后的标签请新建一个labels文件夹，保存在该文件夹下。

生成的xml2yolo.py脚本内容如下：

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

def convert_coordinates(size, box):
    """
    将XML中的边界框坐标转换为YOLO格式
    """
    dw = 1.0/size[0]
    dh = 1.0/size[1]
    
    # XML格式为 xmin, ymin, xmax, ymax
    x = (box[0] + box[2])/2.0
    y = (box[1] + box[3])/2.0
    w = box[2] - box[0]
    h = box[3] - box[1]
    
    # 归一化
    x = x*dw
    w = w*dw
    y = y*dh
    h = h*dh
    
    return (x,y,w,h)

def convert_xml_to_yolo(xml_path, class_mapping):
    """
    转换单个XML文件到YOLO格式
    """
    tree = ET.parse(xml_path)
    root = tree.getroot()
    
    # 获取图像尺寸
    size = root.find('size')
    width = int(size.find('width').text)
    height = int(size.find('height').text)
    
    # 创建对应的txt文件路径
    txt_path = xml_path.replace('Annotations', 'Labels').replace('.xml', '.txt')
    
    # 确保Labels目录存在
    os.makedirs(os.path.dirname(txt_path), exist_ok=True)
    
    with open(txt_path, 'w') as txt_file:
        # 处理每个目标
        for obj in root.iter('object'):
            # 获取类别名称
            class_name = obj.find('name').text
            
            # 获取类别ID
            if class_name not in class_mapping:
                continue
            class_id = class_mapping[class_name]
            
            # 获取边界框坐标
            xmlbox = obj.find('bndbox')
            xmin = float(xmlbox.find('xmin').text)
            ymin = float(xmlbox.find('ymin').text)
            xmax = float(xmlbox.find('xmax').text)
            ymax = float(xmlbox.find('ymax').text)
            
            # 转换坐标
            bb = convert_coordinates((width,height), (xmin,ymin,xmax,ymax))
            
            # 写入txt文件
            txt_file.write(f"{class_id} {bb[0]:.6f} {bb[1]:.6f} {bb[2]:.6f} {bb[3]:.6f}\n")

def main():
    # 定义类别映射
    class_mapping = {
        'missing_hole': 0,
        'mouse_bite': 1,
        'open_circuit': 2,
        'short': 3,
        'spur': 4,
        'spurious_copper': 5
    }
    
    # 获取所有XML文件
    xml_files = glob.glob('Annotations/*/*.xml')
    
    # 转换每个XML文件
    for xml_file in xml_files:
        try:
            convert_xml_to_yolo(xml_file, class_mapping)
            print(f"成功转换: {xml_file}")
        except Exception as e:
            print(f"转换失败 {xml_file}: {str(e)}")

if __name__ == "__main__":
    main()

1.2 数据集划分

cursor prompt

请创建一个py脚本，对数据进行划分。对于image文件夹下的每一个类别按照8:2进行随机划分。请直接生成包含所有图像绝对路径的train.txt和val.txt文件。并输出每个类别训练集验证集的数量。

生成的split_dataset.py脚本内容如下：

import os
import glob
import random

def split_dataset(image_root='images', train_ratio=0.8, seed=42):
    """
    数据集划分函数
    :param image_root: 图片根目录路径
    :param train_ratio: 训练集比例
    :param seed: 随机种子
    """
    # 设置随机种子保证可重复性
    random.seed(seed)
    
    # 初始化路径列表
    train_paths = []
    val_paths = []
    
    # 获取所有类别目录
    class_dirs = [d for d in glob.glob(os.path.join(image_root, '*')) 
                 if os.path.isdir(d)]
    
    # 初始化统计字典
    class_stats = {}
    
    for class_dir in class_dirs:
        # 获取类别名称
        class_name = os.path.basename(class_dir)
        
        # 获取当前类别所有图片路径
        image_paths = glob.glob(os.path.join(class_dir, '*.*'))
        image_paths = [p for p in image_paths 
                      if p.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp'))]
        
        # 打乱顺序
        random.shuffle(image_paths)
        
        # 计算分割点
        split_idx = int(len(image_paths) * train_ratio)
        
        # 分割数据集
        train = image_paths[:split_idx]
        val = image_paths[split_idx:] if split_idx < len(image_paths) else []
        
        # 转换为绝对路径并添加路径分隔符
        train_paths.extend([os.path.abspath(p) + '\n' for p in train])
        val_paths.extend([os.path.abspath(p) + '\n' for p in val])
        
        # 记录统计信息
        class_stats[class_name] = {
            'total': len(image_paths),
            'train': len(train),
            'val': len(val)
        }
    
    # 写入文件
    with open('train.txt', 'w') as f:
        f.writelines(train_paths)
    
    with open('val.txt', 'w') as f:
        f.writelines(val_paths)
    
    # 新增统计信息输出
    print("\n各类别数据分布：")
    print("{:<15} {:<10} {:<10} {:<10}".format('类别', '总数', '训练集', '验证集'))
    for cls, stat in class_stats.items():
        print("{:<15} {:<10} {:<10} {:<10}".format(
            cls, stat['total'], stat['train'], stat['val']
        ))
    
    # 原有总样本数输出保持不变
    print(f'\n数据集划分完成！\n训练集样本数: {len(train_paths)}\n验证集样本数: {len(val_paths)}')

if __name__ == '__main__':
    # 使用示例（根据实际情况修改路径）
    split_dataset(image_root='images')

1.3 yaml配置文件创建

pcb_detect.yaml具体内容如下：

path: E:\project\YOLOv12\dataset\PCB_DATASET # dataset root dir
train: train.txt # train images (relative to 'path') 118287 images
val: val.txt # val images (relative to 'path') 5000 images
test: # 20288 of 40670 images, submit to https://competitions.codalab.org/competitions/20794

# Classes
names:
  0: Missing_hole
  1: Mouse_bite
  2: Open_circuit
  3: Short
  4: Spur
  5: Spurious_copper

2. 训练验证

train.py训练验证脚本内容如下：

from ultralytics import YOLO

model = YOLO('yolov12n.yaml')

# Train the model
results = model.train(
  data='pcb_detect.yaml',
  epochs=300, 
  batch=4, 
  imgsz=640,
  scale=0.5,  # S:0.9; M:0.9; L:0.9; X:0.9
  mosaic=1.0,
  mixup=0.0,  # S:0.05; M:0.15; L:0.15; X:0.2
  copy_paste=0.1,  # S:0.15; M:0.4; L:0.5; X:0.6
  device="0",
  workers=0,
)

# Evaluate model performance on the validation set
metrics = model.val()