from ultralytics import YOLO

if __name__ == '__main__':
    # 載入模型
    model = YOLO('yolov8n-seg.yaml').load('yolov8n-seg.pt')  # 從 YAML 配置並加載權重

    # 訓練參數設定
    img_size = 640  # 訓練影像大小
    num_epochs = 1  # 訓練次數

    # 訓練模型
    print("訓練開始")
    model.train(
        data='E:\\code\\image_segmentation\\YOLODataset\\dataset.yaml',  # 資料集的配置文件
        imgsz=img_size,  # 設置影像大小
        epochs=num_epochs  # 訓練的 epoch 次數
    )
    print("訓練結束")