first commit

app/stream_worker.py

@@ -0,0 +1,221 @@
+from __future__ import annotations
+
+import threading
+import time
+from typing import Any, Callable
+
+import cv2
+
+
+class StreamWorker:
+    def __init__(
+        self,
+        stream_url: str | Callable[[], str],
+        detector: Any,
+        frame_skip: int = 3,
+        jpeg_quality: int = 80,
+        resize_width: int | None = None,
+    ):
+        self.stream_url = stream_url
+        self.detector = detector
+        self.frame_skip = max(1, frame_skip)
+        self.jpeg_quality = jpeg_quality
+        self.resize_width = resize_width
+
+        self.lock = threading.Lock()
+        self.latest_jpeg: bytes | None = None
+        self.latest_detections: list[dict[str, Any]] = []
+        self.frame_id = 0
+        self.updated_at = 0.0
+        self.running = False
+        self.connected = False
+        self.error = "尚未启动"
+        self.fps = 0.0
+        self.thread: threading.Thread | None = None
+        self.reconnect_requested = False
+        self.reconnect_version = 0
+        self.resolve_ms = 0.0
+        self.open_ms = 0.0
+        self.first_frame_ms = 0.0
+
+    def start(self) -> None:
+        if self.running:
+            return
+
+        self.running = True
+        self.thread = threading.Thread(target=self._run, daemon=True)
+        self.thread.start()
+
+    def stop(self) -> None:
+        self.running = False
+        if self.thread and self.thread.is_alive():
+            self.thread.join(timeout=2)
+
+    def reconnect(self) -> None:
+        with self.lock:
+            self.latest_jpeg = None
+            self.latest_detections = []
+            self.frame_id = 0
+            self.fps = 0.0
+            self.reconnect_requested = True
+            self.reconnect_version += 1
+            self.connected = False
+            self.error = "正在切换视频源"
+            self.resolve_ms = 0.0
+            self.open_ms = 0.0
+            self.first_frame_ms = 0.0
+
+    def get_jpeg(self) -> bytes | None:
+        with self.lock:
+            return self.latest_jpeg
+
+    def get_snapshot(self) -> dict[str, Any]:
+        with self.lock:
+            return {
+                "frame_id": self.frame_id,
+                "updated_at": self.updated_at,
+                "detections": list(self.latest_detections),
+                "running": self.running,
+                "connected": self.connected,
+                "error": self.error,
+                "fps": round(self.fps, 2),
+                "timings": {
+                    "resolve_ms": self.resolve_ms,
+                    "open_ms": self.open_ms,
+                    "first_frame_ms": self.first_frame_ms,
+                },
+            }
+
+    def _run(self) -> None:
+        cap: cv2.VideoCapture | None = None
+        last_detections: list[dict[str, Any]] = []
+        fps_window_start = time.monotonic()
+        fps_frames = 0
+
+        while self.running:
+            with self.lock:
+                should_reconnect = self.reconnect_requested
+                run_version = self.reconnect_version
+                self.reconnect_requested = False
+            if should_reconnect:
+                # 切换摄像头时必须释放旧连接，否则 OpenCV 会继续阻塞读旧流。
+                if cap is not None:
+                    cap.release()
+                    cap = None
+
+            if cap is None or not cap.isOpened():
+                started = time.monotonic()
+                stream_url = self.stream_url() if callable(self.stream_url) else self.stream_url
+                resolve_ms = round((time.monotonic() - started) * 1000, 2)
+                started = time.monotonic()
+                cap = cv2.VideoCapture(stream_url)
+                open_ms = round((time.monotonic() - started) * 1000, 2)
+                with self.lock:
+                    self.open_ms = open_ms
+                    self.resolve_ms = resolve_ms
+                    self.first_frame_ms = 0.0
+                if not cap.isOpened():
+                    self._set_connection_state(False, "无法打开视频流，2 秒后重试")
+                    cap.release()
+                    cap = None
+                    time.sleep(2)
+                    continue
+                self._set_connection_state(True, "已连接")
+
+            started = time.monotonic()
+            ok, frame = cap.read()
+            with self.lock:
+                current_version = self.reconnect_version
+            # 丢弃切换期间从旧连接读到的帧，避免前端画面回跳。
+            if current_version != run_version:
+                continue
+            if self.first_frame_ms == 0.0:
+                with self.lock:
+                    self.first_frame_ms = round((time.monotonic() - started) * 1000, 2)
+            if not ok:
+                self._set_connection_state(False, "读取视频帧失败，正在重连")
+                cap.release()
+                cap = None
+                time.sleep(2)
+                continue
+
+            frame = self._resize(frame)
+            self.frame_id += 1
+
+            if self.frame_id % self.frame_skip == 0:
+                frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+                last_detections = self.detector.detect(frame_rgb)
+
+            annotated = self._draw(frame, last_detections)
+            ok, jpeg = cv2.imencode(
+                ".jpg",
+                annotated,
+                [cv2.IMWRITE_JPEG_QUALITY, self.jpeg_quality],
+            )
+            if not ok:
+                continue
+
+            fps_frames += 1
+            now = time.monotonic()
+            if now - fps_window_start >= 1:
+                fps = fps_frames / (now - fps_window_start)
+                fps_window_start = now
+                fps_frames = 0
+            else:
+                fps = self.fps
+
+            with self.lock:
+                current_version = self.reconnect_version
+            if current_version != run_version:
+                continue
+
+            with self.lock:
+                self.latest_jpeg = jpeg.tobytes()
+                self.latest_detections = list(last_detections)
+                self.updated_at = time.time()
+                self.connected = True
+                self.error = ""
+                self.fps = fps
+
+        if cap is not None:
+            cap.release()
+        self._set_connection_state(False, "已停止")
+
+    def _resize(self, frame: Any) -> Any:
+        if not self.resize_width:
+            return frame
+
+        height, width = frame.shape[:2]
+        if width <= self.resize_width:
+            return frame
+
+        scale = self.resize_width / width
+        return cv2.resize(frame, (self.resize_width, int(height * scale)))
+
+    def _draw(self, frame: Any, detections: list[dict[str, Any]]) -> Any:
+        for detection in detections:
+            x1, y1, x2, y2 = detection["box"]
+            label = detection["label"]
+            score = detection["score"]
+
+            cv2.rectangle(frame, (x1, y1), (x2, y2), (0, 220, 80), 2)
+            text = f"{label} {score:.2f}"
+            cv2.rectangle(frame, (x1, max(0, y1 - 26)), (x1 + 150, y1), (0, 220, 80), -1)
+            cv2.putText(
+                frame,
+                text,
+                (x1 + 5, max(18, y1 - 7)),
+                cv2.FONT_HERSHEY_SIMPLEX,
+                0.55,
+                (0, 0, 0),
+                2,
+                cv2.LINE_AA,
+            )
+
+        return frame
+
+    def _set_connection_state(self, connected: bool, error: str) -> None:
+        with self.lock:
+            self.connected = connected
+            self.error = error
+            self.updated_at = time.time()