#!/usr/bin/env python3
"""
YOLOv11 training script for ForestPersons Winter subset
"""
import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
os.environ['MKL_THREADING_LAYER'] = 'GNU'

from ultralytics import YOLO

def train_winter_model():
    """Train YOLOv11 on ForestPersons Winter subset"""
    
    project_name = "forestpersons_winter"  # Project name for winter training results
    
    model_type_list = [
        "yolov11n",
        # "yolov11s",  # Uncomment for larger models
        # "yolov11m",
        # "yolov11l",
    ]
    
    # device = [5, 6]  # GPU devices to use
    device = [4, 5]  # GPU devices to use
    
    for model_type in model_type_list:
        print(f"❄️ Training {model_type} on ForestPersons Winter subset...")
        
        # Load a pretrained YOLO model
        model = YOLO(f"{model_type}.pt")
        
        # Train the model with default hyperparameters
        results = model.train(
            project=project_name,  # Set the project name for saving results
            name=model_type,       # Run name within the project
            device=device,         # Specify the device(s) for training
            data="data_forestpersons_v3/forestpersons_winter.yaml",  # Path to Winter dataset config
        )
        
        print(f"✅ Completed training {model_type} on Winter subset")
        print(f"📊 Results saved in: runs/{project_name}/{model_type}/")
        print("-" * 50)

if __name__ == "__main__":
    print("=" * 60)
    print("❄️ STARTING FORESTPERSONS WINTER TRAINING")
    print("=" * 60)
    
    train_winter_model()
    
    print("\n" + "=" * 60)
    print("🎉 WINTER TRAINING COMPLETED!")
    print("=" * 60)