hookvision-rk/hookvision/rtsp_player.py

import av
import sys
import os
import socket
import struct
import numpy as np
import time
import cv2
import shutil
import subprocess
import threading
import logging
import platform
from urllib.parse import urlparse, urlunparse, quote
from PySide6.QtCore import QObject, Signal, QMutex
from PySide6.QtGui import QImage, QPixmap
from typing import List, Dict, Any, Optional

logger = logging.getLogger(__name__)


def _try_av_open(rtsp_url: str, options: dict):
    """用 av 库打开 RTSP，返回 (container, used_url) 或 (None, None)"""
    urls = [rtsp_url]
    parsed = urlparse(rtsp_url)
    if parsed.username or parsed.password:
        user = quote(parsed.username or '', safe='')
        passwd = quote(parsed.password or '', safe='')
        netloc = f"{user}:{passwd}@{parsed.hostname}"
        if parsed.port:
            netloc += f":{parsed.port}"
        encoded = urlunparse(parsed._replace(netloc=netloc))
        if encoded != rtsp_url:
            urls.append(encoded)

    for url in urls:
        try:
            container = av.open(url, options=options)
            return container, url
        except Exception as e:
            if '401' in str(e) or 'Unauthorized' in str(e):
                logger.warning(f"av 库 RTSP 认证失败 ({url}): {e}")
                continue
            raise
    return None, None


def _find_system_ffmpeg():
    """查找系统 ffmpeg 二进制路径（绕过 PyInstaller 打包的版本）"""
    # 优先用 /usr/bin 下的系统版本
    for path in ['/usr/bin/ffmpeg', '/usr/local/bin/ffmpeg']:
        if os.path.isfile(path) and os.access(path, os.X_OK):
            return path
    # fallback：查找 PATH 中的 ffmpeg
    p = shutil.which('ffmpeg')
    if p:
        return p
    return None


class _FFmpegSubprocess:
    """通过 subprocess 调用系统 ffmpeg 读取 RTSP 流，支持 Digest 认证"""

    def __init__(self, rtsp_url: str, width: int = 1920, height: int = 1080, timeout: float = 10.0):
        self.rtsp_url = rtsp_url
        self.width = width
        self.height = height
        self.timeout = timeout
        self.proc = None
        self._stderr_thread = None
        self._stderr_log = []

    def _read_stderr(self):
        """后台线程读取 ffmpeg stderr，记录错误"""
        try:
            for line in iter(self.proc.stderr.readline, b''):
                line_str = line.decode('utf-8', errors='replace').strip()
                if line_str:
                    self._stderr_log.append(line_str)
                    logger.error(f"[ffmpeg stderr] {line_str}")
        except Exception:
            pass

    def open(self) -> bool:
        ffmpeg_bin = _find_system_ffmpeg()
        if not ffmpeg_bin:
            logger.error("未找到系统 ffmpeg 二进制")
            return False

        cmd = [
            ffmpeg_bin,
            '-hide_banner',
            '-rtsp_transport', 'tcp',
            '-i', self.rtsp_url,
            '-f', 'rawvideo',
            '-pix_fmt', 'rgb24',
            '-vf', f'scale={self.width}:{self.height}',
            '-r', '25',
            'pipe:1'
        ]
        logger.info(f"ffmpeg subprocess 命令: {' '.join(cmd)}")
        try:
            self.proc = subprocess.Popen(
                cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                bufsize=self.width * self.height * 3
            )
        except Exception as e:
            logger.error(f"ffmpeg subprocess 启动失败: {e}")
            return False

        # 后台读取 stderr
        self._stderr_thread = threading.Thread(target=self._read_stderr, daemon=True)
        self._stderr_thread.start()

        # 等待第一帧数据（带超时）
        frame_size = self.width * self.height * 3
        import select
        try:
            ready = select.select([self.proc.stdout], [], [], self.timeout)
            if not ready[0]:
                logger.error(f"ffmpeg subprocess 超时 ({self.timeout}s) 未输出帧数据")
                # 打印 stderr 中的错误
                if self._stderr_log:
                    logger.error(f"ffmpeg 错误输出: {'; '.join(self._stderr_log)}")
                self.proc.kill()
                self.proc = None
                return False
            data = self.proc.stdout.read(frame_size)
            if not data or len(data) < frame_size:
                logger.error(f"ffmpeg subprocess 首帧数据不足: {len(data) if data else 0}/{frame_size}")
                if self._stderr_log:
                    logger.error(f"ffmpeg 错误输出: {'; '.join(self._stderr_log)}")
                self.proc.kill()
                self.proc = None
                return False
            logger.info(f"ffmpeg subprocess 连接成功: {self.rtsp_url}")
            return True
        except Exception as e:
            logger.error(f"ffmpeg subprocess 读取失败: {e}")
            if self._stderr_log:
                logger.error(f"ffmpeg 错误输出: {'; '.join(self._stderr_log)}")
            if self.proc:
                self.proc.kill()
                self.proc = None
            return False

    def read_frame(self):
        if not self.proc or not self.proc.stdout:
            return None
        frame_size = self.width * self.height * 3
        data = self.proc.stdout.read(frame_size)
        if not data or len(data) < frame_size:
            logger.warning(f"ffmpeg subprocess 帧数据不足: {len(data) if data else 0}/{frame_size}")
            return None
        return np.frombuffer(data, dtype=np.uint8).reshape((self.height, self.width, 3))

    def close(self):
        if self.proc:
            try:
                self.proc.kill()
            except Exception:
                pass
            self.proc = None

class RTSPPlayer(QObject):
    frame_ready = Signal(QPixmap)
    first_frame_received = Signal()
    timeout = Signal()

    def __init__(self, rtsp_url: str, target_width: int = 640, target_height: int = 360, timeout_sec: float = 5.0):
        super().__init__()
        self.rtsp_url = rtsp_url
        self.target_width = target_width
        self.target_height = target_height
        self.timeout_sec = timeout_sec
        self.running = False
        self._thread: Optional[threading.Thread] = None
        self._stop_socket: Optional[socket.socket] = None
        self.ai_objects: List[Dict[str, Any]] = []
        self.target_fps = 25
        self.mutex = QMutex()
        self._last_frame_time: float = 0.0
        self._timeout_emitted = False

    def set_target_size(self, width: int, height: int):
        self.mutex.lock()
        self.target_width = width
        self.target_height = height
        self.mutex.unlock()

    def set_ai_objects(self, objects):
        self.mutex.lock()
        self.ai_objects = list(objects)
        self.mutex.unlock()

    def _build_rtsp_urls(self, url: str) -> list:
        """构建原始 URL 和 URL 编码版本，用于认证重试"""
        urls = [url]
        try:
            parsed = urlparse(url)
            if parsed.username or parsed.password:
                user = quote(parsed.username or '', safe='')
                passwd = quote(parsed.password or '', safe='')
                netloc = f"{user}:{passwd}@{parsed.hostname}"
                if parsed.port:
                    netloc += f":{parsed.port}"
                encoded = urlunparse(parsed._replace(netloc=netloc))
                if encoded != url:
                    urls.append(encoded)
        except Exception:
            pass
        return urls

    def run(self):
        self.running = True
        options = {
            'rtsp_transport': 'tcp',
            'buffer_size': '1024',
            'max_delay': '50000',
            'flags': 'low_delay',
            'fflags': 'nobuffer+discardcorrupt',
            'analyzeduration': '50000',
            'probesize': '50000',
            'timeout': str(int(self.timeout_sec * 1000000)),
            'stimeout': str(int(self.timeout_sec * 1000000)),
        }
        start_time = time.time()
        first_packet = True
        container = None
        ffsub = None
        use_ffmpeg_subprocess = False

        # 先尝试 av 库
        try:
            container, used_url = _try_av_open(self.rtsp_url, options)
            if container:
                if used_url != self.rtsp_url:
                    logger.info(f"av 库认证成功（使用编码URL）: {used_url}")
                else:
                    logger.info(f"av 库连接成功: {self.rtsp_url}")
        except Exception as e:
            logger.warning(f"av 库异常: {e}")

        # av 失败时尝试系统 ffmpeg subprocess（支持 Digest 认证）
        if container is None:
            logger.info("av 库失败，尝试系统 ffmpeg subprocess...")
            ffsub = _FFmpegSubprocess(self.rtsp_url, self.target_width, self.target_height, self.timeout_sec)
            if ffsub.open():
                use_ffmpeg_subprocess = True
            else:
                ffsub = None
                logger.error(f"所有方式均失败: {self.rtsp_url}")
                self.timeout.emit()
                self.running = False
                return

        try:
            if use_ffmpeg_subprocess:
                # ===== 系统 ffmpeg subprocess 模式 =====
                self.first_frame_received.emit()
                last_emit = 0
                frame_interval = 1.0 / self.target_fps
                while self.running:
                    img = ffsub.read_frame()
                    if img is None:
                        logger.warning("ffmpeg subprocess 读取帧失败，触发重连")
                        break
                    now = time.time()
                    if now - last_emit >= frame_interval:
                        self.process_frame(img)
                        last_emit = now
                    self._last_frame_time = time.time()
                self._timeout_emitted = True
                self.timeout.emit()
            else:
                # ===== av 库模式 =====
                video_stream = container.streams.video[0]
                video_stream.thread_type = 'AUTO'
                video_stream.thread_count = 1
                last_emit = 0
                frame_interval = 1.0 / self.target_fps

                for packet in container.demux(video_stream):
                    if not self.running:
                        break
                    if first_packet:
                        first_packet = False
                        logger.info(f"RTSP 首包到达耗时 {time.time() - start_time:.1f}s: {self.rtsp_url}")
                        self.first_frame_received.emit()
                    for frame in packet.decode():
                        if not self.running:
                            break
                        img = frame.to_rgb().to_ndarray()
                        now = time.time()
                        if now - last_emit >= frame_interval:
                            self.process_frame(img)
                            last_emit = now
                        self._last_frame_time = time.time()
                    if self._last_frame_time and (time.time() - self._last_frame_time) > self.timeout_sec:
                        logger.warning("RTSP帧超时，触发重连")
                        break
                    if first_packet and (time.time() - start_time) > self.timeout_sec:
                        logger.warning("RTSP连接超时")
                        break
                    if not self.running:
                        break
                self._timeout_emitted = True
                self.timeout.emit()

        except Exception as e:
            logger.error(f"RTSP错误: {e}")
            if not self._timeout_emitted:
                self._timeout_emitted = True
                self.timeout.emit()
        finally:
            if container:
                container.close()
            if ffsub:
                ffsub.close()
            self.running = False

    def start(self):
        self._thread = threading.Thread(target=self.run, daemon=True)
        self._thread.start()

    def stop(self):
        self.running = False

    def is_alive(self):
        return self._thread is not None and self._thread.is_alive()

    def process_frame(self, frame):
        self.mutex.lock()
        tw, th = self.target_width, self.target_height
        objects = list(self.ai_objects)
        self.mutex.unlock()

        orig_h, orig_w = frame.shape[:2]
        if tw > 0 and th > 0 and (orig_w != tw or orig_h != th):
            # 保持原始宽高比缩放到 target 内，不足补黑边
            scale = min(tw / orig_w, th / orig_h)
            new_w = int(orig_w * scale)
            new_h = int(orig_h * scale)
            resized = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
            canvas = np.zeros((th, tw, 3), dtype=np.uint8)
            x_off = (tw - new_w) // 2
            y_off = (th - new_h) // 2
            canvas[y_off:y_off+new_h, x_off:x_off+new_w] = resized
            frame = canvas

            # AI 坐标从 display 空间映射到实际画布内图像区域
            sx = new_w / tw
            sy = new_h / th
        else:
            sx = sy = 1.0
            x_off = y_off = 0

        for obj in objects:
            if obj.get('text'):
                pos = obj.get('pos', (8, 20))
                color = obj.get('color', (0, 255, 0))
                if obj.get('anchor') == 'tr':
                    w = cv2.getTextSize(obj['text'], cv2.FONT_HERSHEY_SIMPLEX, 0.6, 2)[0][0]
                    tx = frame.shape[1] - w - int(pos[0])
                    ty = int(pos[1])
                else:
                    tx = int(pos[0] * sx + x_off)
                    ty = int(pos[1] * sy + y_off)
                cv2.putText(frame, obj['text'], (tx, ty),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 2)
                continue
            if obj.get('points'):
                pts = [(int((p[0] * sx) + x_off), int((p[1] * sy) + y_off)) for p in obj['points']]
                color = obj.get('color', (0, 255, 0))
                alpha = obj.get('alpha', 1.0)
                thickness = obj.get('thickness', 2)
                label = obj.get('label', '')
                if len(pts) >= 2:
                    if alpha < 1.0:
                        overlay = frame.copy()
                        cv2.polylines(overlay, [np.array(pts, dtype=np.int32)], False, color, thickness)
                        cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0, frame)
                    else:
                        cv2.polylines(frame, [np.array(pts, dtype=np.int32)], False, color, thickness)
                if label and pts:
                    tx, ty = pts[-1]
                    cv2.putText(frame, label, (tx + 3, ty - 3), cv2.FONT_HERSHEY_SIMPLEX, 0.4, color, 1)
                continue
            x = int(obj.get('x', 0) * sx + x_off)
            y = int(obj.get('y', 0) * sy + y_off)
            w = int(obj.get('w', 0) * sx)
            h = int(obj.get('h', 0) * sy)
            label = obj.get('label', '')
            color = obj.get('color', (0, 255, 0))
            alpha = obj.get('alpha', 1.0)
            thickness = 1 if alpha < 1.0 else (1 if obj.get('dimmed') else 2)
            if obj.get('dimmed') and alpha >= 1.0:
                color = tuple(c // 2 for c in color)
            if alpha < 1.0:
                overlay = frame.copy()
                cv2.rectangle(overlay, (x, y), (x+w, y+h), color, thickness)
                cv2.putText(overlay, label, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, thickness)
                cv2.addWeighted(overlay, alpha, frame, 1 - alpha, 0, frame)
            else:
                cv2.rectangle(frame, (x, y), (x+w, y+h), color, thickness)
                cv2.putText(frame, label, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, thickness)

        h, w, ch = frame.shape
        bytes_per_line = ch * w
        qt_img = QImage(frame.data, w, h, bytes_per_line, QImage.Format_RGB888).copy()
        pixmap = QPixmap.fromImage(qt_img)
        self.frame_ready.emit(pixmap)