feeding

vision/overflow_detector.py

@@ -1,47 +1,33 @@
 # vision/overflow_detector.py
 import sys
 import os
-from vision.resize_tuili_image_main import classify_image_weighted, load_global_rois, crop_and_resize
+from typing import Optional
+from vision.overflow_model.yiliao_main_rknn import classify_frame_with_rois
 
 # 添加项目根目录到Python路径
-sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
+# sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
 
-def detect_overflow_from_image(image_array, overflow_model, roi_file_path):
+def detect_overflow_from_image(overflow_model,roi_file_path,image_array)->Optional[str]:
     """
     通过图像检测是否溢料
     :param image_array: 图像数组
     :param overflow_model: 溢料检测模型
     :param roi_file_path: ROI文件路径
-    :return: 是否检测到溢料 (True/False)
+    :return: 检测到的溢料类别 (未堆料、小堆料、大堆料、未浇筑满、浇筑满) 或 None
     """
     try:
-        # 检查模型是否已加载
-        if overflow_model is None:
-            print("堆料检测模型未加载")
-            return False
+        outputs = classify_frame_with_rois(overflow_model, image_array, roi_file_path)
+        print("溢料检测结果:", outputs)
+        for res in outputs:
 
-        # 加载ROI区域
-        rois = load_global_rois(roi_file_path)
-
-        if not rois:
-            print(f"没有有效的ROI配置: {roi_file_path}")
-            return False
-
-        if image_array is None:
-            print("输入图像为空")
-            return False
-
-        # 裁剪和调整图像大小
-        crops = crop_and_resize(image_array, rois, 640)
-
-        # 对每个ROI区域进行分类检测
-        for roi_resized, _ in crops:
-            final_class, _, _, _ = classify_image_weighted(roi_resized, overflow_model, threshold=0.4)
-            if "大堆料" in final_class or "小堆料" in final_class:
-                print(f"检测到溢料: {final_class}")
-                return True
-
-        return False
+            return res["class"]
+            # if "大堆料" in res["class"] or "小堆料" in res["class"]:
+            #     print(f"检测到溢料: {res['class']}")
+            #     return True
+        
+        # return False
+        return None
     except Exception as e:
         print(f"溢料检测失败: {e}")
-        return False
+        return None
+