from ultralytics import YOLO

# Caminho para o ficheiro dataset.yaml criado anteriormente
dataset_yaml_path = 'dataset/yolo_format/dataset.yaml'

# Escolhe o modelo base (n = nano, s = small, m = medium, l = large, x = extra)
model_path = 'yolo11s-seg.pt'  # ou yolov8n-seg.pt, yolov8m-seg.pt, etc.

# Carregar modelo
model = YOLO(model_path)

# Treinar o modelo
model.train(
    data=dataset_yaml_path,
    epochs=100,
    imgsz=640,
    batch=24,
    name='meu_yolo_seg',
    project='runs/segment',
    save=True,
    save_period=10,  # Salva a cada 10 épocas
    device=[0,1],
)