121 lines
3.5 KiB
Python
121 lines
3.5 KiB
Python
|
import os
|
|||
|
|
import cv2
|
|||
|
|
from rknnlite.api import RKNNLite
|
|||
|
|
|
|||
|
|
# classify_single_image, StableClassJudge, CLASS_NAMES 已在 muju_cls_rknn 中定义
|
|||
|
|
from muju_cls_rknn import classify_single_image, StableClassJudge, CLASS_NAMES
|
|||
|
|
|
|||
|
|
|
|||
|
|
def run_stable_classification_loop(
|
|||
|
|
model_path,
|
|||
|
|
roi_file,
|
|||
|
|
image_source,
|
|||
|
|
stable_frames=3,
|
|||
|
|
display_scale=0.5, # 显示缩放比例（0.5 = 显示为原来 50%）
|
|||
|
|
show_window=False # 是否显示窗口
|
|||
|
|
):
|
|||
|
|
"""
|
|||
|
|
image_source: cv2.VideoCapture 对象
|
|||
|
|
"""
|
|||
|
|
|
|||
|
|
judge = StableClassJudge(
|
|||
|
|
stable_frames=stable_frames,
|
|||
|
|
ignore_class=2 # 忽略“有遮挡”类别参与稳定判断
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
cap = image_source
|
|||
|
|
if not hasattr(cap, "read"):
|
|||
|
|
raise TypeError("image_source 必须是 cv2.VideoCapture 实例")
|
|||
|
|
|
|||
|
|
# 可选：创建可缩放窗口
|
|||
|
|
if show_window:
|
|||
|
|
cv2.namedWindow("RTSP Stream - Press 'q' to quit", cv2.WINDOW_NORMAL)
|
|||
|
|
|
|||
|
|
while True:
|
|||
|
|
ret, frame = cap.read()
|
|||
|
|
if not ret:
|
|||
|
|
print("无法读取视频帧（可能是流断开或结束）")
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
# 上下左右翻转
|
|||
|
|
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}")
|
|||
|
|
|
|||
|
|
# ---------------------------
|
|||
|
|
# 稳定判断
|
|||
|
|
# ---------------------------
|
|||
|
|
stable_class_id = judge.update(class_id)
|
|||
|
|
|
|||
|
|
if stable_class_id is not None:
|
|||
|
|
print(f"\n稳定输出: {CLASS_NAMES[stable_class_id]}\n")
|
|||
|
|
|
|||
|
|
# ---------------------------
|
|||
|
|
# 显示画面（缩小窗口）
|
|||
|
|
# ---------------------------
|
|||
|
|
if show_window:
|
|||
|
|
h, w = frame.shape[:2]
|
|||
|
|
display_frame = cv2.resize(
|
|||
|
|
frame,
|
|||
|
|
(int(w * display_scale), int(h * display_scale)),
|
|||
|
|
interpolation=cv2.INTER_AREA
|
|||
|
|
)
|
|||
|
|
|
|||
|
|
cv2.imshow("RTSP Stream - Press 'q' to quit", display_frame)
|
|||
|
|
|
|||
|
|
if cv2.waitKey(1) & 0xFF == ord('q'):
|
|||
|
|
break
|
|||
|
|
|
|||
|
|
cap.release()
|
|||
|
|
cv2.destroyAllWindows()
|
|||
|
|
|
|||
|
|
|
|||
|
|
if __name__ == "__main__":
|
|||
|
|
# ---------------------------
|
|||
|
|
# 配置参数
|
|||
|
|
# ---------------------------
|
|||
|
|
MODEL_PATH = "muju_cls.rknn"
|
|||
|
|
ROI_FILE = "./roi_coordinates/muju_roi.txt"
|
|||
|
|
RTSP_URL = "rtsp://admin:XJ123456@192.168.250.61:554/streaming/channels/101"
|
|||
|
|
|
|||
|
|
STABLE_FRAMES = 3
|
|||
|
|
DISPLAY_SCALE = 0.5 # 显示窗口缩放比例
|
|||
|
|
SHOW_WINDOW = False # 部署时改成 False，测试的时候打开
|
|||
|
|
|
|||
|
|
# ---------------------------
|
|||
|
|
# 打开 RTSP 视频流
|
|||
|
|
# ---------------------------
|
|||
|
|
print(f"正在连接 RTSP 流: {RTSP_URL}")
|
|||
|
|
cap = cv2.VideoCapture(RTSP_URL)
|
|||
|
|
|
|||
|
|
# 降低 RTSP 延迟（部分摄像头支持）
|
|||
|
|
cap.set(cv2.CAP_PROP_BUFFERSIZE, 1)
|
|||
|
|
|
|||
|
|
if not cap.isOpened():
|
|||
|
|
print("无法打开 RTSP 流，请检查网络、账号密码或 URL")
|
|||
|
|
exit(1)
|
|||
|
|
|
|||
|
|
print("RTSP 流连接成功，开始推理...")
|
|||
|
|
|
|||
|
|
# ---------------------------
|
|||
|
|
# 启动稳定分类循环三帧稳定判断
|
|||
|
|
# ---------------------------
|
|||
|
|
run_stable_classification_loop(
|
|||
|
|
model_path=MODEL_PATH,
|
|||
|
|
roi_file=ROI_FILE,
|
|||
|
|
image_source=cap,
|
|||
|
|
stable_frames=STABLE_FRAMES,
|
|||
|
|
display_scale=DISPLAY_SCALE,
|
|||
|
|
show_window=SHOW_WINDOW
|
|||
|
|
)
|
|||
|
|
|