from ultralytics import YOLO

def main():
    model = YOLO("models/yolo11n-seg.pt")
    model.train(
        data="./datasets/edges/data.yaml",
        epochs=150,
        patience=20,
        imgsz=640,
        batch=12,
        save_period=10,
        project="result",
        name="edges-nano",
        exist_ok=False,
        device = 0
    )

if __name__ == "__main__":
    main()