video/core/rtmp.py

import asyncio
import logging
import cv2
import time
from ocr.model_violation_detector import MultiModelViolationDetector


# 配置文件相对路径（根据实际目录结构调整）
YOLO_MODEL_PATH = "../ocr/models/best.pt"  # 关键修正：从core目录向上一级找ocr文件夹
FORBIDDEN_WORDS_PATH = "../ocr/forbidden_words.txt"
OCR_CONFIG_PATH = "../ocr/config/1.yaml"
KNOWN_FACES_DIR = "../ocr/known_faces"

# 创建检测器实例
detector = MultiModelViolationDetector(
    forbidden_words_path=FORBIDDEN_WORDS_PATH,
    ocr_config_path=OCR_CONFIG_PATH,
    yolo_model_path=YOLO_MODEL_PATH,
    known_faces_dir=KNOWN_FACES_DIR,
    ocr_confidence_threshold=0.5
)

# 配置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger("rtmp_video_puller")


async def rtmp_pull_video_stream(rtmp_url):
    """
    通过RTMP从指定URL拉取视频流并进行违规检测
    """
    cap = None  # 初始化视频捕获对象
    try:
        # 异步打开RTMP流
        cap = await asyncio.to_thread(
            cv2.VideoCapture,
            rtmp_url,
            cv2.CAP_FFMPEG  # 指定FFmpeg后端确保RTMP兼容性
        )

        # 检查RTMP流是否成功打开
        is_opened = await asyncio.to_thread(cap.isOpened)
        if not is_opened:
            raise Exception(f"RTMP流打开失败: {rtmp_url}（请检查URL有效性和FFmpeg环境）")

        # 获取RTMP流基础信息
        width = await asyncio.to_thread(cap.get, cv2.CAP_PROP_FRAME_WIDTH)
        height = await asyncio.to_thread(cap.get, cv2.CAP_PROP_FRAME_HEIGHT)
        fps = await asyncio.to_thread(cap.get, cv2.CAP_PROP_FPS)

        # 处理异常情况
        fps = fps if fps > 0 else 30.0
        width, height = int(width), int(height)

        # 打印流初始化成功信息
        print(f"RTMP流状态: 已成功连接")
        print(f"流基础信息: 分辨率 {width}x{height} | 配置帧率 {fps:.2f} FPS")
        print("开始接收视频帧...（按 Ctrl+C 中断）")

        # 初始化帧统计参数
        frame_count = 0
        start_time = time.time()

        # 循环读取视频帧
        while True:
            ret, frame = await asyncio.to_thread(cap.read)

            if not ret:
                print(f"RTMP流状态: 帧读取失败（可能流已中断或结束）")
                break

            frame_count += 1

            # 打印当前帧信息
            print(f"收到帧 (第{frame_count}帧)")
            print(f"  帧尺寸: {width}x{height}")
            print(f"  配置帧率: {fps:.2f} FPS")

            if frame is not None:
                has_violation, violation_type, details = detector.detect_violations(frame)
                if has_violation:
                    print(f"检测到违规 - 类型: {violation_type}, 详情: {details}")
                else:
                    print("未检测到任何违规内容")
            else:
                print(f"无法读取测试图像")

            # 每100帧统计一次实际接收帧率
            if frame_count % 100 == 0:
                elapsed_time = time.time() - start_time
                actual_fps = frame_count / elapsed_time
                print(f"---- 帧统计: 累计{frame_count}帧 | 实际平均帧率 {actual_fps:.2f} FPS ----")

    except KeyboardInterrupt:
        print(f"\n用户操作: 已通过 Ctrl+C 中断程序")
    except Exception as e:
        logger.error(f"RTMP流处理异常: {str(e)}", exc_info=True)
        print(f"错误信息: {str(e)}")
    finally:
        if cap is not None:
            await asyncio.to_thread(cap.release)
            print(f"\n资源释放: RTMP流已关闭")
            print(f"最终统计: 共接收 {frame_count if 'frame_count' in locals() else 0} 帧")


if __name__ == "__main__":
    RTMP_URL = "rtmp://192.168.110.25:1935/live/473b95a47e338301cbd96809ea7ac416"

    try:
        asyncio.run(rtmp_pull_video_stream(RTMP_URL))
    except Exception as e:
        print(f"程序启动失败: {str(e)}")