57 lines
1.9 KiB
Python
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}")
|