from ultralytics import YOLO

# Load model
model = YOLO("yolo26m.pt")

# Train with added augmentation
model.train(
    data="dataset.yaml", 
    imgsz=640, 
    batch=16, 
    epochs=50,
    patience=10,
    #fraction=0.05, # currently uses 100% of the dataset; 0.8 = 80% of the set
    save_period=5, 
    workers=0, 
    device=0,
    #close_mosaic=10,

    # Augmentation settings
    augment=True,      # General augmentation on/off
    degrees=10,        # Rotation up to 10 degrees
    translate=0.1,     # Translation up to 10%
    scale=0.5,         # Scaling between 0.5 and 1.5
    fliplr=0.5,        # 50% chance of horizontal flip
    hsv_h=0.015,       # Minor hue variations
    hsv_s=0.7,         # Saturation variations
    hsv_v=0.4,         # Brightness variations
    mosaic=1.0         # Mosaic augmentation (combines 4 images)
)