AP/Yolo_predict.py

57 lines
1.9 KiB
Python

import os
import cv2
import datetime
import threading
from ultralytics import YOLO
class YoloPredict(threading.Thread):
""" YOLO 推論執行緒 """
def __init__(self, image, model_path, save_dir="results", callback=None):
super(YoloPredict, self).__init__()
self.image = image
self.model_path = model_path
self.save_dir = save_dir
self.callback = callback # ✅ 設定回呼函數
self.model = None
# ✅ 確保 `results/` 資料夾存在
if not os.path.exists(self.save_dir):
os.makedirs(self.save_dir)
def run(self):
""" 執行 YOLO 推論 """
if self.image is None:
print("⚠️ 無影像可進行推論")
return
try:
# ✅ 載入 YOLO 模型
if self.model is None:
self.model = YOLO(self.model_path)
print("✅ YOLO 模型載入成功")
# ✅ 轉換影像格式 (BGR → RGB)
image_rgb = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
# ✅ 使用 YOLO 模型進行推論
results = self.model.predict(image_rgb, imgsz=640, conf=0.5)
print("✅ YOLO 推論完成")
# ✅ 取得標註結果
result_image = results[0].plot()
# ✅ 儲存影像
timestamp = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
save_path = os.path.join(self.save_dir, f"{timestamp}.jpg")
result_image_bgr = cv2.cvtColor(result_image, cv2.COLOR_RGB2BGR)
cv2.imwrite(save_path, result_image_bgr)
print(f"✅ 推論結果儲存至: {save_path}")
# ✅ 回傳結果到主視窗 (`callback` 函數)
if self.callback:
self.callback(result_image, save_path)
except Exception as e:
print(f"⚠️ 推論時發生錯誤: {e}")