27 lines
972 B
Python
27 lines
972 B
Python
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("訓練結束")
|