78 lines
2.4 KiB
Python
78 lines
2.4 KiB
Python
import cv2
|
|
from detect_bag import detect_bag
|
|
|
|
THRESHOLD_X = 537 # min_x 阈值
|
|
|
|
def bag_judgment(img, return_conf=True, return_vis=False):
|
|
"""
|
|
判断图片中的料袋状态，可动态返回置信度和可视化图像
|
|
Args:
|
|
img (np.ndarray): 待检测图片
|
|
return_conf (bool): 是否返回置信度
|
|
return_vis (bool): 是否返回可视化图像
|
|
Returns:
|
|
status_bool: True=到位, False=未到位, None=未检测到
|
|
status_text: 中文状态
|
|
conf: 最大置信度或 None
|
|
min_x: 最左边 x 坐标或 None
|
|
vis_img: 可视化图像或 None
|
|
"""
|
|
# 调用 detect_bag
|
|
outputs = detect_bag(img, return_conf=return_conf, return_vis=return_vis)
|
|
|
|
# 初始化占位
|
|
conf = None
|
|
min_x = None
|
|
vis_img = None
|
|
|
|
# 根据返回值长度解析
|
|
if return_conf and return_vis:
|
|
if len(outputs) == 3:
|
|
conf, min_x, vis_img = outputs
|
|
elif len(outputs) == 2:
|
|
conf, min_x = outputs
|
|
elif len(outputs) == 1:
|
|
min_x = outputs[0]
|
|
elif return_conf and not return_vis:
|
|
if len(outputs) >= 2:
|
|
conf, min_x = outputs[:2]
|
|
elif len(outputs) == 1:
|
|
min_x = outputs[0]
|
|
elif not return_conf and return_vis:
|
|
if len(outputs) == 2:
|
|
min_x, vis_img = outputs
|
|
elif len(outputs) == 1:
|
|
min_x = outputs[0]
|
|
else:
|
|
min_x = outputs if isinstance(outputs, (int, float, np.number)) else outputs[0]
|
|
|
|
# 判断状态
|
|
if min_x is None:
|
|
status_bool = None
|
|
status_text = "没有料袋"
|
|
elif min_x >= THRESHOLD_X:
|
|
status_bool = True
|
|
status_text = "料袋到位"
|
|
else:
|
|
status_bool = False
|
|
status_text = "料袋未到位"
|
|
|
|
return status_bool, status_text, conf, min_x, vis_img
|
|
|
|
|
|
# ====================== 测试 ======================
|
|
if __name__ == "__main__":
|
|
IMG_PATH = "3.jpg"
|
|
img = cv2.imread(IMG_PATH)
|
|
if img is None:
|
|
raise FileNotFoundError(f"图片无法读取: {IMG_PATH}")
|
|
|
|
status_bool, status_text, conf, min_x, vis_img = bag_judgment(img, return_conf=True, return_vis=True)
|
|
print(f"判断结果: {status_bool}, 中文状态: {status_text}, conf={conf}, min_x={min_x}")
|
|
|
|
if vis_img is not None:
|
|
cv2.imshow("Vis", vis_img)
|
|
cv2.waitKey(0)
|
|
cv2.destroyAllWindows()
|
|
|