224 lines
7.4 KiB
Python
224 lines
7.4 KiB
Python
"""视频流工作线程：读取视频帧、调用检测器、绘制 OSD 并缓存 JPEG。"""
|
||
|
||
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()
|