#!/usr/bin/env python3
"""
YOLOv11 training script for ForestPersons Summer subset
"""

import os
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
os.environ['MKL_THREADING_LAYER'] = 'GNU'

from ultralytics import YOLO

def train_summer_model():
    """Train YOLOv11 on ForestPersons Summer subset"""
    
    project_name = "forestpersons_summer_0924"  # Project name for summer training results
    
    model_type_list = [
        "yolov11n",
        # "yolov11s",  # Uncomment for larger models
        # "yolov11m",
        # "yolov11l",
    ]
    
    # device = [5, 6]  # GPU devices to use
    device = [1,2]  # GPU devices to use
    
    for model_type in model_type_list:
        print(f"🌞 Training {model_type} on ForestPersons Summer 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_summer.yaml",  # Path to Summer dataset config
        )
        
        print(f"✅ Completed training {model_type} on Summer subset")
        print(f"📊 Results saved in: runs/{project_name}/{model_type}/")
        print("-" * 50)

if __name__ == "__main__":
    print("=" * 60)
    print("🌞 STARTING FORESTPERSONS SUMMER TRAINING")
    print("=" * 60)
    
    train_summer_model()
    
    print("\n" + "=" * 60)
    print("🎉 SUMMER TRAINING COMPLETED!")
    print("=" * 60)