image_segmentation/train_seg.py
2025-01-20 16:21:14 +08:00

19 lines
565 B
Python

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("訓練結束")