import os
from ultralytics import YOLO
import argparse

if __name__ == '__main__':
    # 參數解析
    parser = argparse.ArgumentParser(description='Train YOLOv8 model')
    parser.add_argument('--data-dir', type=str, default='./YOLODataset', help='Directory containing dataset configuration file')
    parser.add_argument('--img-size', type=int, default=640, help='Training image size')
    parser.add_argument('--epochs', type=int, default=1, help='Number of training epochs')
    args = parser.parse_args()

    # 路徑設定
    data_yaml_path = os.path.join(args.data_dir, 'dataset.yaml')

    # 載入模型
    model = YOLO('yolov8s.yaml').load('yolov8s.pt')  # 從 YAML 配置並加載權重

    # 訓練模型
    print("訓練開始")
    model.train(
        data=data_yaml_path,  # 動態獲取資料集的配置文件
        imgsz=args.img_size,  # 設置影像大小
        epochs=args.epochs  # 訓練的 epoch 次數
    )
    print("訓練結束")