from ultralytics import YOLO
from tkinter import Tk, filedialog

# 載入模型
model = YOLO('E:/code/image_segmentation/runs/segment/train14/weights/best.pt')

# 使用 tkinter 開啟檔案選擇對話框
def select_files():
    root = Tk()
    root.withdraw()  # 隱藏主視窗
    file_paths = filedialog.askopenfilenames(
        title="Select Images",
        filetypes=[("Image files", "*.bmp;*.png;*.jpg;*.jpeg")]
    )
    return file_paths

# 選擇影像檔案
image_paths = select_files()
if not image_paths:
    print("No files selected.")
else:
    # 執行推論
    for image_path in image_paths:
        results = model(
            source=image_path,  # 輸入圖片路徑
            save=True,  # 儲存推論結果
            device='0',  # 使用 GPU 若發生錯誤改成CPU
            # max_det=8,  # 每個類別只保留一個檢測結果
            # iou=0.1,  # 降低 IOU 閾值，進一步減少重複檢測
            # conf=0.7  # 可以根據需要調整置信度閾值
        )
    print("Inference completed!")