from ultralytics import YOLO
from tkinter import Tk, filedialog
import cv2
import numpy as np
import os

# 載入模型
model = YOLO('./runs/detect/train3/weights/best.pt')

# 設定輸出資料夾
results_dir = "./results"
os.makedirs(results_dir, exist_ok=True)  # 確保結果資料夾存在


# 使用 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
            conf=0.4  # 可以根據需要調整信心度開關
        )

        # 傾向只顯示方框和信心度
        image = cv2.imread(image_path)
        for r in results:
            if r.boxes is not None:
                for box in r.boxes:
                    if box.conf.item() > 0.4:  # 過濾信心度
                        x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
                        confidence = box.conf.item()

                        # 顯示方框
                        cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
                        cv2.putText(image, f'{confidence:.2f}', (x1, y1 - 10),
                                    cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 0), 2)

        # 設定輸出路徑到 results 資料夾
        output_filename = os.path.basename(image_path).replace(".jpg", "_output.jpg").replace(".png", "_output.png")
        output_path = os.path.join(results_dir, output_filename)
        cv2.imwrite(output_path, image)
        print(f"Saved result to {output_path}")

    print("Inference completed!")