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

# 載入模型
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,
            device='0',  # 使用 GPU 若發生錯誤改成CPU
            max_det=8,
            iou=0.1,
            conf=0.7
        )

        # 取得圖片名稱和目錄
        img_name = os.path.basename(image_path)
        output_dir = os.path.join(os.path.dirname(image_path), "segmented_results")
        os.makedirs(output_dir, exist_ok=True)

        # 處理每張圖片的結果
        for result in results:
            original_image = result.orig_img
            masks = result.masks.data.cpu().numpy() if result.masks else None

            if masks is not None:
                # 初始化遮罩影像
                mask_image = np.zeros_like(original_image)

                for mask in masks:
                    mask = mask.astype(np.uint8) * 255

                    # 調整遮罩大小與原始影像一致
                    mask_resized = cv2.resize(mask, (original_image.shape[1], original_image.shape[0]))

                    # 將單通道遮罩轉換為三通道
                    mask_colored = cv2.merge([mask_resized, mask_resized, mask_resized])

                    # 疊加遮罩到影像
                    mask_image = cv2.addWeighted(mask_image, 1, mask_colored, 1, 0)

                # 保存結果
                output_path = os.path.join(output_dir, f"segmented_{img_name}")
                cv2.imwrite(output_path, mask_image)
                print(f"Saved segmented result to: {output_path}")
    print("Inference completed!")