34 lines
1.1 KiB
Python
34 lines
1.1 KiB
Python
# vision/overflow_detector.py
|
|
import sys
|
|
import os
|
|
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__), '..'))
|
|
|
|
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: 检测到的溢料类别 (未堆料、小堆料、大堆料、未浇筑满、浇筑满) 或 None
|
|
"""
|
|
try:
|
|
outputs = classify_frame_with_rois(overflow_model, image_array, roi_file_path)
|
|
print("溢料检测结果:", outputs)
|
|
for res in outputs:
|
|
|
|
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 None
|
|
|