1. 数据集准备

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

1.1 标签格式转换

cursor prompt

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

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

import os
import xml.etree.ElementTree as ET
import globdef convert_coordinates(size, box):"""将XML中的边界框坐标转换为YOLO格式"""dw = 1.0/size[0]dh = 1.0/size[1]# XML格式为 xmin, ymin, xmax, ymaxx = (box[0] + box[2])/2.0y = (box[1] + box[3])/2.0w = box[2] - box[0]h = box[3] - box[1]# 归一化x = x*dww = w*dwy = y*dhh = h*dhreturn (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# 获取类别IDif class_name not in class_mapping:continueclass_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 randomdef 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_hole1: Mouse_bite2: Open_circuit3: Short4: Spur5: Spurious_copper

2. 训练验证

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

from ultralytics import YOLOmodel = 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.9mosaic=1.0,mixup=0.0,  # S:0.05; M:0.15; L:0.15; X:0.2copy_paste=0.1,  # S:0.15; M:0.4; L:0.5; X:0.6device="0",workers=0,
)# Evaluate model performance on the validation set
metrics = model.val()

遇到``AttributeError: ‘InfiniteDataLoader‘ object has no attribute ‘` 报错，查看解决方案~