61 lines
2.0 KiB
Python
61 lines
2.0 KiB
Python
from ultralytics import YOLO
|
|
from tkinter import Tk, filedialog
|
|
import cv2
|
|
import numpy as np
|
|
import os
|
|
|
|
# 載入模型
|
|
model = YOLO('./runs/detect/train3/weights/best.pt')
|
|
|
|
# 設定輸出資料夾
|
|
results_dir = "./results"
|
|
os.makedirs(results_dir, exist_ok=True) # 確保結果資料夾存在
|
|
|
|
|
|
# 使用 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, # 輸入圖片路徑
|
|
save=True, # 儲存推論結果
|
|
device='0', # 使用 GPU 若發生錯誤改成CPU
|
|
conf=0.4 # 可以根據需要調整信心度開關
|
|
)
|
|
|
|
# 傾向只顯示方框和信心度
|
|
image = cv2.imread(image_path)
|
|
for r in results:
|
|
if r.boxes is not None:
|
|
for box in r.boxes:
|
|
if box.conf.item() > 0.4: # 過濾信心度
|
|
x1, y1, x2, y2 = map(int, box.xyxy[0].tolist())
|
|
confidence = box.conf.item()
|
|
|
|
# 顯示方框
|
|
cv2.rectangle(image, (x1, y1), (x2, y2), (255, 0, 0), 2)
|
|
cv2.putText(image, f'{confidence:.2f}', (x1, y1 - 10),
|
|
cv2.FONT_HERSHEY_SIMPLEX, 2, (255, 0, 0), 2)
|
|
|
|
# 設定輸出路徑到 results 資料夾
|
|
output_filename = os.path.basename(image_path).replace(".jpg", "_output.jpg").replace(".png", "_output.png")
|
|
output_path = os.path.join(results_dir, output_filename)
|
|
cv2.imwrite(output_path, image)
|
|
print(f"Saved result to {output_path}")
|
|
|
|
print("Inference completed!")
|