# YOLOv8 segmentation training for car damage detection
from ultralytics import YOLO
import multiprocessing
import os

def train():
    # Start from YOLOv8 medium segmentation model
    model = YOLO('../../models/yolov8m-seg.pt')

    # Get the absolute path to the data.yaml file
    current_dir = os.path.dirname(os.path.abspath(__file__))
    data_yaml_path = os.path.join(current_dir, 'data.yaml')

    # Train with optimized parameters
    model.train(
        data=data_yaml_path,  # Path to data configuration file
        epochs=150,  # Number of epochs
        imgsz=640,  # Image size
        batch=4,    # Batch size
        workers=4,  # Number of workers
        project='../../models/damage/weights',  # Save directory
        name='yolov8_damage_final',  # Run name
        
        # Learning rate strategy
        lr0=0.0002,  # Initial learning rate
        lrf=0.000001,  # Final learning rate
        warmup_epochs=25,
        warmup_momentum=0.8,
        cos_lr=True,  # Use cosine learning rate scheduler
        
        # Loss weights
        box=8.0,  # Box loss gain
        cls=4.0,  # Class loss gain
        dfl=2.5,  # DFL loss gain
        
        # Augmentation settings
        augment=True,
        mosaic=0.5,
        mixup=0.2,
        copy_paste=0.1,
        degrees=20.0,
        translate=0.2,
        scale=0.4,
        shear=10.0,
        flipud=0.1,
        fliplr=0.5,
        hsv_h=0.015,
        hsv_s=0.7,
        hsv_v=0.4,
        
        # Other optimization settings
        overlap_mask=True,  # Overlap mask segments
        mask_ratio=4,  # Mask downsampling ratio
        single_cls=True,  # Single class detection
        rect=False,  # Rectangular training
        cache=False,  # Cache images for faster training
        patience=50,  # Early stopping patience
        close_mosaic=10,  # Close mosaic augmentation epochs
        deterministic=True,  # Deterministic mode
        seed=42,  # Random seed
        device=0  # GPU device
    )

if __name__ == '__main__':
    multiprocessing.freeze_support()
    train()