code-test-drone/ground_station/video_source.py

"""video_source.py —— 吊舱 RTSP 视频源（服务端拉流 → MJPEG 推流）。

浏览器不能直连 RTSP，由本模块在服务端用 OpenCV 拉流、编码 JPEG 帧缓存，
页面经 `/video_feed`（multipart/x-mixed-replace）显示实时画面。

设计要点：
  * cv2 延迟导入：模块 import 无副作用，无 OpenCV 环境也能 import；
    capture_factory 可注入假采集器 → 单测不依赖真实摄像头/网络；
  * 断线自动重连（读帧失败 → 释放 → 退避重试）；
  * 抽帧节流 + 降分辨率：1080p60 全量转 JPEG 带宽/CPU 吃不消（默认 960 宽、≤15fps）。
"""

import logging
import threading
import time

logger = logging.getLogger(__name__)

_IDLE_RESEND_S = 2.0   # 无新帧时重发当前帧作保活（浏览器挂住连接）


class RtspVideoSource:
    """后台线程拉 RTSP，维护最新 JPEG 帧供 MJPEG 消费。"""

    def __init__(self, url: str, width: int = 960, fps: float = 15.0,
                 quality: int = 75, capture_factory=None) -> None:
        self._url = url
        self._width = int(width)
        self._min_interval = 1.0 / max(1.0, float(fps))
        self._quality = int(quality)
        self._capture_factory = capture_factory
        self._lock = threading.Lock()
        self._frame = None            # 最新 JPEG bytes
        self._frame_seq = 0
        self._frame_at = 0.0          # 最新帧时间（monotonic）
        self._stop = threading.Event()
        self._thread = None

    # --- 生命周期 -----------------------------------------------------------
    def start(self) -> None:
        if self._thread is not None and self._thread.is_alive():
            return
        self._stop.clear()
        self._thread = threading.Thread(target=self._run, name="rtsp-video",
                                        daemon=True)
        self._thread.start()

    def stop(self) -> None:
        self._stop.set()
        thread, self._thread = self._thread, None
        if thread is not None:
            thread.join(timeout=3.0)

    @property
    def online(self) -> bool:
        """最近 2s 内有新帧 → 在线（用于页面状态显示）。"""
        with self._lock:
            return (self._frame is not None
                    and (time.monotonic() - self._frame_at) < 2.0)

    def latest(self) -> tuple[int, bytes | None]:
        """返回 (seq, jpeg)；尚无帧时 (0, None)。"""
        with self._lock:
            return self._frame_seq, self._frame

    def mjpeg_stream(self):
        """MJPEG 生成器（multipart 分块）；客户端断开时随 GeneratorExit 终止。"""
        last_seq = -1
        last_send = 0.0
        while not self._stop.is_set():
            seq, jpeg = self.latest()
            now = time.monotonic()
            if jpeg is not None and (seq != last_seq
                                     or now - last_send > _IDLE_RESEND_S):
                last_seq = seq
                last_send = now
                yield (b"--frame\r\nContent-Type: image/jpeg\r\n"
                       + f"Content-Length: {len(jpeg)}\r\n\r\n".encode("ascii")
                       + jpeg + b"\r\n")
            time.sleep(1.0 / 30.0)

    # --- 拉流线程 -----------------------------------------------------------
    def _open(self):
        if self._capture_factory is not None:
            return self._capture_factory()
        import cv2  # 延迟导入：仅真正拉流时依赖 OpenCV
        return cv2.VideoCapture(self._url, cv2.CAP_FFMPEG)

    def _run(self) -> None:
        import os
        # stimeout：RTSP socket 超时（µs）。不设的话断网时单次 open 会阻塞 ~30s，
        # 影响重连节奏（实测 2026-10-14）。
        os.environ.setdefault("OPENCV_FFMPEG_CAPTURE_OPTIONS",
                              "rtsp_transport;tcp|stimeout;5000000")
        import cv2

        cap = None
        last_encode = 0.0
        try:
            while not self._stop.is_set():
                if cap is None or not cap.isOpened():
                    if cap is not None:
                        try:
                            cap.release()
                        except Exception:
                            pass
                        cap = None
                    cap = self._open()
                    if not cap.isOpened():
                        logger.warning("[Video] RTSP 打开失败，1s 后重试: %s",
                                       self._url)
                        self._stop.wait(1.0)
                        continue
                    logger.info("[Video] RTSP 已连接: %s", self._url)
                    last_encode = 0.0

                ok, frame = cap.read()
                if not ok or frame is None:
                    logger.warning("[Video] 读帧失败，重连")
                    try:
                        cap.release()
                    except Exception:
                        pass
                    cap = None
                    self._stop.wait(0.5)
                    continue

                now = time.monotonic()
                if now - last_encode < self._min_interval:
                    continue  # 抽帧节流：60fps 源 → ≤fps 编码
                last_encode = now

                h, w = frame.shape[:2]
                if w > self._width:
                    scale = self._width / w
                    frame = cv2.resize(
                        frame, (self._width, max(1, int(h * scale))),
                        interpolation=cv2.INTER_AREA)
                ok, buf = cv2.imencode(
                    ".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, self._quality])
                if not ok:
                    continue
                with self._lock:
                    self._frame = buf.tobytes()
                    self._frame_seq += 1
                    self._frame_at = time.monotonic()
        finally:
            if cap is not None:
                try:
                    cap.release()
                except Exception:
                    pass