# -*- 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()

train: ../train/images
val: ../valid/images
test: ../test/images
path: ../data/images

nc: 33
names: ['Acalolepta', 'Aromia', 'Batocera rubus', 'Batocera rufomaculata', 
'Blaberus discoidalis', 'Brown marmorated stink bug', 'Caterpillar', 'Cepora nerissa', 
'Cerura Vinula', 'Clostera', 'Crowned slug moth', 'Empoasca', 'Eurydema ornata',
'Fall webworm moth', 'Forest Bug', 'Graphosoma lineatum', 'Neocurtilla hexadactyla', 
'Oebalus pugnax', 'Papilio memnon', 'Parasa consocia', 'Parasa indetermina', 'Passionvine hopper',
'Plagiodera versicolora', 'Protaetia', 'Puss moth', 'Red Costate Tiger Moth', 'Spilomelinae', 
'Spirama', 'Spotted lanternfly', 'Tibicina haematodes', 'Torpedo bug', 'White-spotted longicorn beetle', 
'Yellow spotted stink bug']