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 segmentation model
    model = YOLO("yolov8l-seg.pt")  # Using large model as medium might be too small for segmentation

    # Training configuration
    results = model.train(
        data=r"C:\Users\trist\Documents\School\2024-2025\Semester_2\Project one\2024-2025-projectone-ctai-DebrabandereTristan\AI\Datasets\DartboardSegmentation.v19i.yolov8\data.yaml",
        
        # Essential parameters
        epochs=50,
        imgsz=640,
        patience=10,
        batch=8 if device != "cpu" else 4,
        device=device,
        name="yolov8l_seg_darts",
        workers=4,  # Reduced from 8 to prevent potential memory issues
        
        # Optimization parameters
        optimizer="AdamW",
        lr0=0.001,
        lrf=0.01,
        momentum=0.937,
        weight_decay=0.0005,
        warmup_epochs=3.0,
        warmup_momentum=0.8,
        warmup_bias_lr=0.1,
        
        # Augmentation parameters
        hsv_h=0.015,
        hsv_s=0.7,
        hsv_v=0.4,
        degrees=45.0,
        translate=0.1,
        scale=0.5,
        shear=0.0,
        perspective=0.0001,
        flipud=0.5,
        fliplr=0.5,
        mosaic=1.0,
        mixup=0.1,
        copy_paste=0.1,
        auto_augment="randaugment",
        
        # Segmentation-specific parameters
        mask_ratio=4,
        overlap_mask=True,
        
        # Regularization
        dropout=0.1,
        
        # Validation settings
        val=True,
        save_period=10,
        plots=True,
        close_mosaic=10,
        
        # Advanced settings
        deterministic=False,
        erasing=0.4,
    )