89 lines
2.9 KiB
Python
89 lines
2.9 KiB
Python
import os
|
|
import cv2
|
|
#from rknnlite.api import RKNNLite
|
|
import time
|
|
|
|
# classify_single_image, StableClassJudge, CLASS_NAMES 已在 muju_cls_rknn 中定义
|
|
from .charge_cls_rknn import classify_single_image, StableClassJudge, CLASS_NAMES
|
|
|
|
# 获取当前文件所在目录的绝对路径
|
|
current_dir = os.path.dirname(os.path.abspath(__file__))
|
|
def run_stable_charge_loop():
|
|
"""
|
|
image_source: cv2.VideoCapture 对象
|
|
"""
|
|
_ret=None
|
|
# 使用相对于当前文件的绝对路径
|
|
model_path = os.path.join(current_dir, "charge0324.rknn")
|
|
# roi_file = os.path.join(current_dir, "roi_coordinates", "muju_roi.txt")
|
|
RTSP_URL = "rtsp://admin:XJ123456@192.168.250.60:554/streaming/channels/101"
|
|
stable_frames=5
|
|
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)
|
|
|
|
|
|
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=="插好":
|
|
break
|
|
print(f"-------当前振捣棒检测为：{_ret}，继续等待稳定------")
|
|
else:
|
|
print("-------当前振捣棒检测为空，继续等待稳定------")
|
|
|
|
time.sleep(0.1)
|
|
finally:
|
|
if cap is not None:
|
|
cap.release()
|
|
return _ret
|
|
|
|
|
|
|
|
|