# -*- coding: utf-8 -*-
"""
該腳本用於執行YOLO模型的訓練。

它會自動處理以下任務：
1. 動態修改數據集配置文件 (data.yaml)，將相對路徑更新為絕對路徑，以確保訓練時能正確找到數據。
2. 從 'pretrained' 文件夾加載指定的預訓練模型。
3. 使用預設的參數（如epochs, imgsz, batch）啓動訓練過程。

要開始訓練，只需直接運行此腳本。
"""
import os
import yaml
from pathlib import Path
from ultralytics import YOLO

def main():
    """
    主訓練函數。
    
    該函數負責執行YOLO模型的訓練流程，包括：
    1. 配置預訓練模型。
    2. 動態修改數據集的YAML配置文件，確保路徑為絕對路徑。
    3. 加載預訓練模型。
    4. 使用指定參數開始訓練。
    """
    # --- 1. 配置模型和路徑 ---
    
    # 要訓練的模型列表
    models_to_train = [
        {'name': 'yolov5nu.pt', 'train_name': 'train_yolov5nu'},
        {'name': 'yolov8n.pt', 'train_name': 'train_yolov8n'},
        {'name': 'yolo11n.pt', 'train_name': 'train_yolo11n'},
        {'name': 'yolo12n.pt', 'train_name': 'train_yolo12n'}
    ]
    
    # 獲取當前工作目錄的絕對路徑，以避免相對路徑帶來的問題
    current_dir = os.path.abspath(os.getcwd())
    
    # --- 2. 動態配置數據集YAML文件 ---
    
    # 構建數據集yaml文件的絕對路徑
    data_yaml_path = os.path.join(current_dir, 'train_data', 'data.yaml')
    
    # 讀取原始yaml文件內容
    with open(data_yaml_path, 'r', encoding='utf-8') as f:
        data_config = yaml.safe_load(f)
    
    # 將yaml文件中的 'path' 字段修改為數據集目錄的絕對路徑
    # 這是為了確保ultralytics庫能正確定位到訓練、驗證和測試集
    data_config['path'] = os.path.join(current_dir, 'train_data')
    
    # 將修改後的配置寫回yaml文件
    with open(data_yaml_path, 'w', encoding='utf-8') as f:
        yaml.dump(data_config, f, default_flow_style=False, allow_unicode=True)
    
    # --- 3. 循環訓練每個模型 ---
    
    for model_info in models_to_train:
        model_name = model_info['name']
        train_name = model_info['train_name']
        
        print(f"\n{'='*60}")
        print(f"開始訓練模型: {model_name}")
        print(f"訓練名稱: {train_name}")
        print(f"{'='*60}")
        
        # 構建預訓練模型的完整路徑
        pretrained_model_path = os.path.join(current_dir, 'pretrained', model_name)
        if not os.path.exists(pretrained_model_path):
            print(f"警告: 預訓練模型文件不存在: {pretrained_model_path}")
            print(f"跳過模型 {model_name} 的訓練")
            continue
        
        try:
            # 加載指定的預訓練模型
            model = YOLO(pretrained_model_path)
            
            # --- 4. 開始訓練 ---
            
            print(f"開始訓練 {model_name}...")
            # 調用train方法開始訓練
            model.train(
                data=data_yaml_path,  # 數據集配置文件
                epochs=100,           # 訓練輪次
                imgsz=640,            # 輸入圖像尺寸
                batch=8,             # 每批次的圖像數量
                name=train_name,      # 模型名稱
            )
            
            print(f"{model_name} 訓練完成！")
            
        except Exception as e:
            print(f"訓練 {model_name} 時出現錯誤: {str(e)}")
            print(f"跳過模型 {model_name}，繼續訓練下一個模型")
            continue
    
    print(f"\n{'='*60}")
    print("所有模型訓練完成！")
    print(f"{'='*60}")

if __name__ == "__main__":
    # 當該腳本被直接執行時，調用main函數
    main()

# YOLO數據集配置文件
path: labels  # 數據集根目錄
train: train/images  # 訓練集圖像路徑
val: valid/images      # 驗證集圖像路徑
path: data


# 類別數量
nc: 18

# 類別名稱
names: ['yoke', 'yoke suspension', 'spacer', 'stockbridge damper', 'lightning rod shackle', 'lightning rod suspension', 'polymer insulator', 'glass insulator', 'tower id plate', 'vari-grip', 'polymer insulator lower shackle', 'polymer insulator upper shackle', 'polymer insulator tower shackle', 'glass insulator big shackle', 'glass insulator small shackle', 'glass insulator tower shackle', 'spiral damper', 'sphere']