110 lines
4.4 KiB
Python
110 lines
4.4 KiB
Python
|
import cv2
|
||
|
|
from ultralytics import YOLO
|
||
|
|
import numpy as np
|
||
|
|
from PyQt5.QtWidgets import QApplication, QFileDialog, QListView, QTreeView, QAbstractItemView
|
||
|
|
import os
|
||
|
|
import shutil
|
||
|
|
from collections import defaultdict
|
||
|
|
|
||
|
|
|
||
|
|
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 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}")
|
||
|
|
|
||
|
|
|
||
|
|
def process_and_move_images(yolo, input_dir, error_base_dir):
|
||
|
|
moved_count = defaultdict(int)
|
||
|
|
expected_class = os.path.basename(input_dir)
|
||
|
|
|
||
|
|
for filename in os.listdir(input_dir):
|
||
|
|
if filename.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
|
||
|
|
image_path = os.path.join(input_dir, filename)
|
||
|
|
image = cv2.imread(image_path)
|
||
|
|
results = yolo.YoloDetect(image)
|
||
|
|
result_Max_Index = results[0].probs.top1 # 獲取分類編號
|
||
|
|
confidence = results[0].probs.top1conf.item() # 獲取 top1 的置信度
|
||
|
|
|
||
|
|
# 獲取類型名稱
|
||
|
|
predicted_class = get_class_name(result_Max_Index)
|
||
|
|
print(f"原始圖片: {image_path}")
|
||
|
|
print(f"推論結果: {predicted_class}, 信心值: {confidence}")
|
||
|
|
|
||
|
|
# 如果預測類別與期望類別不同且信心值高於 0.4，則移動圖片
|
||
|
|
if predicted_class != expected_class and confidence > 0.4:
|
||
|
|
# 只創建預測類別資料夾
|
||
|
|
error_class_dir = os.path.join(error_base_dir, predicted_class)
|
||
|
|
if not os.path.exists(error_class_dir):
|
||
|
|
os.makedirs(error_class_dir)
|
||
|
|
|
||
|
|
new_image_path = os.path.join(error_class_dir, filename)
|
||
|
|
shutil.move(image_path, new_image_path)
|
||
|
|
print(f"移動圖片至: {new_image_path}")
|
||
|
|
moved_count[predicted_class] += 1
|
||
|
|
else:
|
||
|
|
print("圖片未移動")
|
||
|
|
print("---")
|
||
|
|
|
||
|
|
total_moved = sum(moved_count.values())
|
||
|
|
print(f"總共移動了 {total_moved} 張圖片")
|
||
|
|
print("各類型移動數量：")
|
||
|
|
for class_name, count in moved_count.items():
|
||
|
|
print(f"{class_name}: {count} 張")
|
||
|
|
|
||
|
|
return moved_count
|
||
|
|
|
||
|
|
|
||
|
|
def process_multiple_folders(yolo, input_dirs):
|
||
|
|
for input_dir in input_dirs:
|
||
|
|
print(f"處理資料夾: {input_dir}")
|
||
|
|
error_base_dir = os.path.join(os.path.dirname(input_dir), "Error")
|
||
|
|
moved_count = process_and_move_images(yolo, input_dir, error_base_dir)
|
||
|
|
print(f"資料夾 {input_dir} 處理完成\n")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == '__main__':
|
||
|
|
app = QApplication([])
|
||
|
|
yolo = yolo_class(r"D:\ScrewdriverFile\train_0811_Milwaukee_41\weights\best.pt")
|
||
|
|
|
||
|
|
# 選擇多個輸入資料夾
|
||
|
|
dialog = QFileDialog()
|
||
|
|
dialog.setFileMode(QFileDialog.DirectoryOnly)
|
||
|
|
dialog.setOption(QFileDialog.DontUseNativeDialog, True)
|
||
|
|
|
||
|
|
file_view = dialog.findChild(QListView, 'listView')
|
||
|
|
if file_view:
|
||
|
|
file_view.setSelectionMode(QAbstractItemView.ExtendedSelection)
|
||
|
|
|
||
|
|
tree_view = dialog.findChild(QTreeView, 'treeView')
|
||
|
|
if tree_view:
|
||
|
|
tree_view.setSelectionMode(QAbstractItemView.ExtendedSelection)
|
||
|
|
|
||
|
|
if dialog.exec_():
|
||
|
|
input_dirs = dialog.selectedFiles()
|
||
|
|
if input_dirs:
|
||
|
|
process_multiple_folders(yolo, input_dirs)
|
||
|
|
print("所有選定的資料夾處理完成")
|
||
|
|
else:
|
||
|
|
print("未選擇資料夾")
|
||
|
|
else:
|
||
|
|
print("取消選擇資料夾")
|