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.ai_objects = objects

    def run(self):
        self.running = True
        options = {
            'rtsp_transport': 'tcp',
            'buffer_size': '1024',
            'max_delay': '50000',
            'flags': 'low_delay',
            'fflags': 'nobuffer+discardcorrupt',
            'analyzeduration': '100000',
            'probesize': '100000',
            '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)
            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
                    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.mutex.lock()
                        tw, th = self.target_width, self.target_height
                        self.mutex.unlock()
                        if tw > 0 and th > 0 and (img.shape[1] != tw or img.shape[0] != th):
                            img = cv2.resize(img, (tw, th), interpolation=cv2.INTER_LINEAR)

                        # 再在缩放后的图像上绘制 AI 标注
                        for obj in self.ai_objects:
                            x = obj.get('x', 0)
                            y = obj.get('y', 0)
                            w = obj.get('w', 0)
                            h = obj.get('h', 0)
                            label = obj.get('label', '')
                            color = obj.get('color', (0, 255, 0))
                            cv2.rectangle(img, (x, y), (x+w, y+h), color, 2)
                            cv2.putText(img, label, (x, y-5), cv2.FONT_HERSHEY_SIMPLEX, 0.5, color, 2)

                        # 转换为 QPixmap
                        h, w, ch = img.shape
                        bytes_per_line = ch * w
                        qt_img = QImage(img.data, w, h, bytes_per_line, QImage.Format_RGB888)
                        pixmap = QPixmap.fromImage(qt_img)
                        self.frame_ready.emit(pixmap)
                        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()