hookvision-rk/hookvision/rtsp_player.py

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

logger = logging.getLogger(__name__)

class RTSPPlayer(QThread):
    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.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 stop(self):
        self.running = False
        self.wait()

    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)