Feeding_control_system/vision/muju_cls/muju_utils.py

import os
import cv2
#from rknnlite.api import RKNNLite
import time

# classify_single_image, StableClassJudge, CLASS_NAMES 已在 muju_cls_rknn 中定义
from .muju_cls_rknn import classify_single_image, StableClassJudge, CLASS_NAMES

# 获取当前文件所在目录的绝对路径
current_dir = os.path.dirname(os.path.abspath(__file__))
def run_stable_classification_loop():
    """
    image_source: cv2.VideoCapture 对象
    """
    _ret=None
    # 使用相对于当前文件的绝对路径
    model_path = os.path.join(current_dir, "muju_cls.rknn")
    roi_file = os.path.join(current_dir, "roi_coordinates", "muju_roi.txt")
    RTSP_URL = "rtsp://admin:XJ123456@192.168.250.61:554/streaming/channels/101"
    stable_frames=3
    print(f"正在连接 RTSP 流: {RTSP_URL}")
    cap =None
    try:
        cap = cv2.VideoCapture(RTSP_URL)
        # 降低 RTSP 延迟（部分摄像头支持）
        cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)

        if not cap.isOpened():
            print("无法打开 RTSP 流，请检查网络、账号密码或 URL")
            return None

        print("RTSP 流连接成功，开始推理...")

        judge = StableClassJudge(
            stable_frames=stable_frames,
            ignore_class=2  # 忽略“有遮挡”类别参与稳定判断
        )

        if not hasattr(cap, "read"):
            raise TypeError("image_source 必须是 cv2.VideoCapture 实例")
        _max_count=10
        while True:
            _max_count=_max_count-1

            ret, frame = cap.read()
            if not ret:
                print("无法读取视频帧（可能是流断开或结束）")
                continue
            # 上下左右翻转
            frame = cv2.flip(frame, -1)

            # ---------------------------
            # 单帧推理
            # ---------------------------
            result = classify_single_image(frame, model_path, roi_file)


            class_id = result["class_id"]
            class_name = result["class"]
            score = result["score"]

            print(f"[FRAME] {class_name} | conf={score:.3f}")
            if score>0.8:
                # ---------------------------
                # 稳定判断
                # ---------------------------
                stable_class_id = judge.update(class_id)

                if stable_class_id is not None:
                    _ret=CLASS_NAMES[stable_class_id]
                    if _ret is None:
                        print("-------当前类别为：为空，继续等待稳定------")
                        continue
                    if _ret=="模具车1" or _ret=="模具车2":
                        break
                    print(f"当前类别为：{_ret}，继续等待稳定...")
                else:
                    print("当前类别为：为空，继续等待稳定...")
                
            time.sleep(0.1)
    finally:
        if cap is not None:
            cap.release()
    return _ret