125 lines
5.2 KiB
Python
125 lines
5.2 KiB
Python
import cv2
|
|
from ultralytics import YOLO
|
|
import numpy as np
|
|
from PyQt5.QtWidgets import QApplication, QFileDialog
|
|
import os
|
|
import shutil
|
|
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\AOI_Origin_Error\Error")
|
|
file_paths, _ = file_dialog.getOpenFileNames(None, "Select Images", "", "Image Files (*.png *.jpg *.bmp)")
|
|
return file_paths
|
|
|
|
|
|
def get_class_name(index):
|
|
# 更新類型名稱對應表，包含 41 個類別
|
|
class_names = {
|
|
0: "HexTamperproof_1", 1: "HexTamperproof_2", 2: "HexTamperproof_3", 3: "HexTamperproof_4",
|
|
4: "Hex_1", 5: "Hex_2", 6: "Hex_3",
|
|
7: "Opposite_1",
|
|
8: "Pentalope_1", 9: "Pentalope_2",
|
|
10: "Phillips_1", 11: "Phillips_2", 12: "Phillips_3", 13: "Phillips_4", 14: "Phillips_5",
|
|
15: "Slotted_1", 16: "Slotted_2", 17: "Slotted_3", 18: "Slotted_4", 19: "Slotted_5", 20: "Slotted_6",
|
|
21: "Spanner_1", 22: "Spanner_2",
|
|
23: "Square_1", 24: "Square_2",
|
|
25: "Standoff_1",
|
|
26: "TORXTamperproof_1", 27: "TORXTamperproof_2", 28: "TORXTamperproof_3", 29: "TORXTamperproof_4",
|
|
30: "TORXTamperproof_5", 31: "TORXTamperproof_6", 32: "TORXTamperproof_7",
|
|
33: "TORX_1", 34: "TORX_2",
|
|
35: "TriWing_1", 36: "TriWing_2", 37: "TriWing_3",
|
|
38: "Triangle_1", 39: "Triangle_2",
|
|
40: "Vacancy_1"
|
|
}
|
|
return class_names.get(index, f"Unknown_{index}")
|
|
|
|
# 0: "phillips", 1: "triwing", 2: "standoff", 3: "spanner", 4: "hex",
|
|
# 5: "hextamperproof", 6: "slotted", 7: "torx", 8: "torxtamperproof",
|
|
# 9: "square", 10: "triangle", 11: "pentalope", 12: "phillips_1",
|
|
# 13: "phillips_2", 14: "phillips_3", 15: "phillips_4", 16: "vacancy",
|
|
# 17: "opposite"
|
|
|
|
# 0: "HexTamperproof_1", 1: "HexTamperproof_2", 2: "HexTamperproof_3", 3: "HexTamperproof_4",
|
|
# 4: "Hex_1", 5: "Hex_2", 6: "Hex_3",
|
|
# 7: "Opposite_1",
|
|
# 8: "Pentalope_1", 9: "Pentalope_2",
|
|
# 10: "Phillips_1", 11: "Phillips_2", 12: "Phillips_3", 13: "Phillips_4", 14: "Phillips_5",
|
|
# 15: "Slotted_1", 16: "Slotted_2", 17: "Slotted_3", 18: "Slotted_4", 19: "Slotted_5", 20: "Slotted_6",
|
|
# 21: "Spanner_1", 22: "Spanner_2",
|
|
# 23: "Square_1", 24: "Square_2",
|
|
# 25: "Standoff_1",
|
|
# 26: "TORXTamperproof_1", 27: "TORXTamperproof_2", 28: "TORXTamperproof_3", 29: "TORXTamperproof_4",
|
|
# 30: "TORXTamperproof_5", 31: "TORXTamperproof_6", 32: "TORXTamperproof_7",
|
|
# 33: "TORX_1", 34: "TORX_2",
|
|
# 35: "TriWing_1", 36: "TriWing_2", 37: "TriWing_3",
|
|
# 38: "Triangle_1", 39: "Triangle_2",
|
|
# 40: "Vacancy_1"
|
|
def process_and_rename_images(yolo, image_paths, base_output_dir):
|
|
for image_path in image_paths:
|
|
image = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
|
|
results = yolo.YoloDetect(image)
|
|
|
|
result_Max_Index = results[0].probs.top1 # 獲取分類編號
|
|
confidence = results[0].probs.top1conf.item() # 獲取 top1 的置信度
|
|
|
|
# 獲取類型名稱
|
|
class_name = get_class_name(result_Max_Index)
|
|
|
|
print(f"原始圖片: {image_path}")
|
|
print(f"推論結果: {class_name}, 信心值: {confidence}")
|
|
|
|
# 只有當信心值高於 0.6 時才保存圖片
|
|
if confidence > 0.4:
|
|
# 獲取當前日期和時間，精度到毫秒
|
|
current_time = datetime.now().strftime("%Y%m%d_%H%M%S_%f")[:-1]
|
|
|
|
# 獲取原始檔案的副檔名
|
|
_, file_extension = os.path.splitext(image_path)
|
|
|
|
# 創建新的檔名，包含日期時間
|
|
new_filename = f"{class_name}_conf{confidence:.8f}{file_extension}"
|
|
|
|
# 創建對應類別的資料夾
|
|
class_output_dir = os.path.join(base_output_dir, class_name)
|
|
if not os.path.exists(class_output_dir):
|
|
os.makedirs(class_output_dir)
|
|
|
|
# 設置新檔案的完整路徑
|
|
new_filepath = os.path.join(class_output_dir, new_filename)
|
|
|
|
# 複製並重命名檔案
|
|
shutil.copy2(image_path, new_filepath)
|
|
|
|
print(f"新檔名: {new_filename}")
|
|
print(f"儲存位置: {class_output_dir}")
|
|
print("儲存成功")
|
|
else:
|
|
print("信心值過低，不儲存圖片")
|
|
|
|
print("---")
|
|
|
|
|
|
if __name__ == '__main__':
|
|
app = QApplication([])
|
|
yolo = yolo_class(r"D:\ScrewdriverFile\train_0811_Milwaukee_41\weights\best.pt")
|
|
# D:\ScrewdriverFile\train_0811_Milwaukee_41\weights\best.pt
|
|
# C:\Users\user\Desktop\0723ProgramBackup\model\1c344897-1978-4664-9bbd-256be14a3125.pt
|
|
image_paths = open_image_files()
|
|
if image_paths:
|
|
base_output_dir = r"D:\Screwdriver_image\AOI_Origin_Error\Error_rename"
|
|
process_and_rename_images(yolo, image_paths, base_output_dir)
|
|
else:
|
|
print("未選擇任何圖片") |