90 lines
2.9 KiB
Python
90 lines
2.9 KiB
Python
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}") |