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 ‘` 報錯，查看解決方案~