48 lines
1.5 KiB
Python
48 lines
1.5 KiB
Python
# vision/overflow_detector.py
|
|
import sys
|
|
import os
|
|
from vision.resize_tuili_image_main import classify_image_weighted, load_global_rois, crop_and_resize
|
|
|
|
# 添加项目根目录到Python路径
|
|
sys.path.append(os.path.join(os.path.dirname(__file__), '..'))
|
|
|
|
def detect_overflow_from_image(image_array, overflow_model, roi_file_path):
|
|
"""
|
|
通过图像检测是否溢料
|
|
:param image_array: 图像数组
|
|
:param overflow_model: 溢料检测模型
|
|
:param roi_file_path: ROI文件路径
|
|
:return: 是否检测到溢料 (True/False)
|
|
"""
|
|
try:
|
|
# 检查模型是否已加载
|
|
if overflow_model is None:
|
|
print("堆料检测模型未加载")
|
|
return False
|
|
|
|
# 加载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
|
|
except Exception as e:
|
|
print(f"溢料检测失败: {e}")
|
|
return False
|