from ultralytics import YOLO
import torch

if __name__ == "__main__":
    # Device configuration
    device = "0" if torch.cuda.is_available() else "cpu"
    print(f"Using device: {'CUDA:' + device if device != 'cpu' else 'CPU'}")

    # Load YOLOv8 pose model
    model = YOLO("yolov8m-pose.pt")  # Pose estimation model
    
    # Training configuration for dart keypoint detection
    results = model.train(
        data=r"C:\Users\trist\Documents\School\2024-2025\Semester_2\Project one\2024-2025-projectone-ctai-DebrabandereTristan\AI\Datasets\dartkey.v3i.yolov8\data.yaml",
        
        # Essential parameters
        epochs=200,
        imgsz=640,
        batch=16 if device != "cpu" else 8,
        device=device,
        name="yolov8m_dart_pose",
        workers=4,
        
        # Optimization
        optimizer="AdamW",
        lr0=0.01,
        lrf=0.01,
        momentum=0.937,
        weight_decay=0.0005,
        warmup_epochs=3.0,
        
        # Augmentations (adjusted for keypoints)
        hsv_h=0.015,
        hsv_s=0.7,
        hsv_v=0.4,
        degrees=10.0,  # Reduced rotation
        translate=0.1,
        scale=0.5,
        shear=0.0,  # Disabled for keypoints
        perspective=0.0005,
        flipud=0.0,  # Disabled vertical flip
        fliplr=0.5,  # Horizontal flip
        mosaic=1.0,
        mixup=0.0,  # Disabled
        copy_paste=0.0,  # Disabled
        
        # Regularization
        dropout=0.1,
        
        # Validation
        val=True,
        save_period=10,
        plots=True,
        patience=50  # Early stopping
    )