import cv2
from ultralytics import YOLO
import numpy as np
from PyQt5.QtWidgets import QApplication, QFileDialog
import os
from datetime import datetime


class yolo_class():
    def __init__(self, yolo_path):
        self.yolo_path = yolo_path
        self.model = YOLO(yolo_path)  # load a custom model

    def YoloDetect(self, imgs):
        results = self.model(imgs)  # predict on an image
        return results


def open_image_files():
    # 打開檔案對話框，允許選擇多個檔案
    file_dialog = QFileDialog()
    file_dialog.setDirectory(r"D:\Screwdriver_image\Milwaukee_0927\AOI_Origin")
    file_dialog.setFileMode(QFileDialog.ExistingFiles)
    file_paths, _ = file_dialog.getOpenFileNames(None, "Select Images", "", "Image Files (*.png *.jpg *.bmp)")
    return file_paths


def save_results(ok_count, ng_count, results_detail):
    # 建立結果資料夾（如果不存在）
    if not os.path.exists('detection_results'):
        os.makedirs('detection_results')

    # 取得現在時間作為檔名
    current_time = datetime.now().strftime("%Y%m%d_%H%M%S")
    result_file = f'detection_results/detection_result_{current_time}.txt'

    with open(result_file, 'w', encoding='utf-8') as f:
        f.write(f"檢測時間: {current_time}\n")
        f.write(f"OK數量: {ok_count}\n")
        f.write(f"NG數量: {ng_count}\n")
        f.write("\n詳細結果:\n")
        for result in results_detail:
            f.write(f"{result}\n")

    return result_file


if __name__ == '__main__':
    app = QApplication([])
    yolo = yolo_class(r"C:\Users\user\Desktop\0723ProgramBackup\model\1c344897-1978-4664-9bbd-256be14a3125.pt")
    # D:\ScrewdriverFile\train_0811_Milwaukee_41\weights\best.pt

    image_paths = open_image_files()
    ok_count = 0
    ng_count = 0
    results_detail = []

    if image_paths:
        for image_path in image_paths:
            image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
            results = yolo.YoloDetect(image)

            result_Max_Index = [result.probs.top1 for result in results]
            confidences = [result.probs.top1conf.item() for result in results]

            # 判斷結果
            is_ng = any(conf < 0.8 for conf in confidences)
            result_str = "NG" if is_ng else "OK"

            # 更新計數
            if is_ng:
                ng_count += 1
            else:
                ok_count += 1

            # 準備詳細結果
            detail = f"圖片: {os.path.basename(image_path)}, 結果: {result_str}"
            detail += f", 推論結果: {result_Max_Index}, 信心值: {confidences}"
            results_detail.append(detail)

            # 印出即時結果
            print(detail)

    # 儲存結果到檔案
    if results_detail:
        result_file = save_results(ok_count, ng_count, results_detail)
        print(f"\n結果總覽:")
        print(f"OK數量: {ok_count}")
        print(f"NG數量: {ng_count}")
        print(f"詳細結果已儲存至: {result_file}")