from ultralytics import YOLO

def main():
    model = YOLO("models/yolo11s.pt")
    model.train(
        data="./datasets/pieces/unified/data.yaml",
        epochs=150,
        patience=20,
        imgsz=640,
        batch=12,
        save_period=10,
        project="result",
        name="unified-small",
        exist_ok=True,
        device=0,
        augment=True,
        flipud=0.5,          # vertical flip with 50% probability
        fliplr=0.5,          # horizontal flip with 50% probability
        hsv_h=0.015,         # change hue
        hsv_s=0.7,           # change saturation
        hsv_v=0.4,           # change brightness/value
        degrees=10.0,        # random rotation ±10 degrees
        translate=0.1,       # translation ±10%
        scale=0.1,           # scaling ±10%
        shear=2.0,           # shear ±2 degrees
        mosaic=1.0,          # mosaic augmentation
        mixup=0.5            # mixup augmentation
    )

if __name__ == "__main__":
    main()