hookvision-rk/hookvision/rtsp_player.py

import av
import numpy as np
import time
import cv2
import threading
import logging
from PySide6.QtCore import QObject, Signal, QMutex
from PySide6.QtGui import QImage, QPixmap
from typing import List, Dict, Any, Optional

logger = logging.getLogger(__name__)

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()

    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 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

        try:
            container = av.open(self.rtsp_url, options=options)
            open_time = time.time() - start_time
            logger.info(f"RTSP av.open 耗时 {open_time:.1f}s: {self.rtsp_url}")
            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

                if first_packet and (time.time() - start_time) > self.timeout_sec:
                    logger.warning("RTSP连接超时（未收到任何包）")
                    self.timeout.emit()
                    self.running = False
                    break

                if not self.running:
                    break

        except Exception as e:
            logger.error(f"RTSP错误: {e}")
            if first_packet:
                self.timeout.emit()
        finally:
            if container:
                container.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:
            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))
            dimmed = obj.get('dimmed', False)
            thickness = 1 if dimmed else 2
            if dimmed:
                color = tuple(c // 2 for c in color)
            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)