202 lines
7.7 KiB
Python
202 lines
7.7 KiB
Python
from __future__ import annotations
|
|
|
|
from dataclasses import dataclass
|
|
from datetime import datetime, timezone
|
|
from pathlib import Path
|
|
from typing import Any, Iterable
|
|
from urllib.parse import urlsplit
|
|
|
|
from .domain import Box, Detection
|
|
|
|
try:
|
|
import cv2 # type: ignore
|
|
import numpy as np # type: ignore
|
|
except ImportError: # pragma: no cover - exercised by the startup failure path
|
|
cv2 = None
|
|
np = None
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class FramePacket:
|
|
frame: Any
|
|
captured_at: datetime
|
|
scripted_detections: tuple[Detection, ...] | None
|
|
|
|
|
|
def require_opencv() -> None:
|
|
if cv2 is None or np is None:
|
|
raise RuntimeError("Brain demo requires the pinned NumPy and OpenCV packages")
|
|
|
|
|
|
def read_stream_url(path_value: str) -> str:
|
|
path = Path(path_value)
|
|
if not path.is_absolute():
|
|
raise ValueError("stream URL file must be an absolute external path")
|
|
try:
|
|
resolved = path.resolve(strict=True)
|
|
except OSError as exc:
|
|
raise ValueError("stream URL file does not exist") from exc
|
|
if not resolved.is_file():
|
|
raise ValueError("stream URL file does not exist")
|
|
repository_root = Path(__file__).resolve().parents[2]
|
|
try:
|
|
resolved.relative_to(repository_root)
|
|
except ValueError:
|
|
pass
|
|
else:
|
|
raise ValueError("stream URL file must be outside the repository")
|
|
try:
|
|
with resolved.open("rb") as stream:
|
|
raw = stream.read(4097)
|
|
except OSError as exc:
|
|
raise ValueError("stream URL file cannot be read") from exc
|
|
if len(raw) > 4096:
|
|
raise ValueError("stream URL file exceeds 4096 bytes")
|
|
try:
|
|
lines = raw.decode("utf-8").splitlines()
|
|
except UnicodeError as exc:
|
|
raise ValueError("stream URL file must be UTF-8") from exc
|
|
values = [line.strip() for line in lines if line.strip()]
|
|
if len(values) != 1:
|
|
raise ValueError("stream URL file must contain exactly one non-empty line")
|
|
parsed = urlsplit(values[0])
|
|
if parsed.scheme not in {"rtsp", "rtsps"} or not parsed.hostname:
|
|
raise ValueError("stream URL must be an RTSP URL")
|
|
return values[0]
|
|
|
|
|
|
class SyntheticSource:
|
|
mode = "synthetic"
|
|
label = "合成回放"
|
|
fixture = True
|
|
source_ref = "demo-camera-01"
|
|
|
|
def __init__(self, width: int = 960, height: int = 540) -> None:
|
|
require_opencv()
|
|
self.width = width
|
|
self.height = height
|
|
self._sequence = 0
|
|
|
|
def read(self) -> FramePacket:
|
|
self._sequence += 1
|
|
frame = np.zeros((self.height, self.width, 3), dtype=np.uint8)
|
|
frame[:] = (20, 28, 42)
|
|
cv2.rectangle(frame, (0, int(self.height * 0.72)), (self.width, self.height), (31, 42, 58), -1)
|
|
for x in range(0, self.width, 80):
|
|
cv2.line(frame, (x, int(self.height * 0.72)), (x + 80, self.height), (43, 57, 75), 1)
|
|
cv2.putText(frame, "SYNTHETIC FIXTURE - NOT MODEL OUTPUT", (24, 38), cv2.FONT_HERSHEY_SIMPLEX, 0.72, (82, 190, 245), 2)
|
|
|
|
phase = ((self._sequence - 1) % 160) / 159.0
|
|
center_x = -0.04 + phase * 1.08
|
|
x1 = max(0.0, center_x - 0.04)
|
|
x2 = min(1.0, center_x + 0.04)
|
|
detections: tuple[Detection, ...] = ()
|
|
if x2 - x1 > 0.01:
|
|
box = Box(x1, 0.34, x2, 0.82)
|
|
detections = (Detection("P-DEMO-001", "person", box, None),)
|
|
px = int(center_x * self.width)
|
|
head_y = int(self.height * 0.40)
|
|
cv2.circle(frame, (px, head_y), 17, (195, 210, 225), -1)
|
|
cv2.line(frame, (px, head_y + 18), (px, int(self.height * 0.65)), (195, 210, 225), 12)
|
|
cv2.line(frame, (px, int(self.height * 0.52)), (px - 30, int(self.height * 0.60)), (195, 210, 225), 8)
|
|
cv2.line(frame, (px, int(self.height * 0.52)), (px + 30, int(self.height * 0.60)), (195, 210, 225), 8)
|
|
cv2.line(frame, (px, int(self.height * 0.65)), (px - 24, int(self.height * 0.79)), (195, 210, 225), 9)
|
|
cv2.line(frame, (px, int(self.height * 0.65)), (px + 24, int(self.height * 0.79)), (195, 210, 225), 9)
|
|
return FramePacket(frame, datetime.now(timezone.utc), detections)
|
|
|
|
def close(self) -> None:
|
|
return
|
|
|
|
|
|
class StreamSource:
|
|
mode = "stream"
|
|
label = "外部 MediaMTX / RTSP"
|
|
fixture = False
|
|
source_ref = "configured-video-source"
|
|
|
|
def __init__(self, stream_url: str) -> None:
|
|
require_opencv()
|
|
self._stream_url = stream_url
|
|
self._capture: Any = None
|
|
|
|
def _open(self) -> None:
|
|
if self._capture is not None:
|
|
self._capture.release()
|
|
self._capture = cv2.VideoCapture(self._stream_url)
|
|
self._capture.set(cv2.CAP_PROP_BUFFERSIZE, 1)
|
|
|
|
def read(self) -> FramePacket:
|
|
if self._capture is None or not self._capture.isOpened():
|
|
self._open()
|
|
ok, frame = self._capture.read()
|
|
if not ok or frame is None:
|
|
self._open()
|
|
raise RuntimeError("stream frame unavailable")
|
|
return FramePacket(frame, datetime.now(timezone.utc), None)
|
|
|
|
def close(self) -> None:
|
|
if self._capture is not None:
|
|
self._capture.release()
|
|
self._capture = None
|
|
|
|
|
|
class HOGPersonDetector:
|
|
name = "opencv_hog_person_demo"
|
|
production_ready = False
|
|
|
|
def __init__(self) -> None:
|
|
require_opencv()
|
|
self._hog = cv2.HOGDescriptor()
|
|
self._hog.setSVMDetector(cv2.HOGDescriptor_getDefaultPeopleDetector())
|
|
|
|
def detect(self, frame: Any) -> list[tuple[Box, float]]:
|
|
height, width = frame.shape[:2]
|
|
scale = min(1.0, 960.0 / max(width, 1))
|
|
working = frame if scale == 1.0 else cv2.resize(frame, (int(width * scale), int(height * scale)))
|
|
boxes, weights = self._hog.detectMultiScale(working, winStride=(8, 8), padding=(8, 8), scale=1.05)
|
|
result: list[tuple[Box, float]] = []
|
|
work_height, work_width = working.shape[:2]
|
|
for raw_box, weight in zip(boxes, weights):
|
|
x, y, box_width, box_height = (int(value) for value in raw_box)
|
|
x1 = max(0.0, min(1.0, x / work_width))
|
|
y1 = max(0.0, min(1.0, y / work_height))
|
|
x2 = max(0.0, min(1.0, (x + box_width) / work_width))
|
|
y2 = max(0.0, min(1.0, (y + box_height) / work_height))
|
|
if x2 > x1 and y2 > y1:
|
|
result.append((Box(x1, y1, x2, y2), float(weight)))
|
|
return result
|
|
|
|
|
|
class CentroidTracker:
|
|
def __init__(self, max_distance: float = 0.18, ttl_frames: int = 8) -> None:
|
|
self._max_distance = max_distance
|
|
self._ttl_frames = ttl_frames
|
|
self._next_id = 1
|
|
self._tracks: dict[str, tuple[Box, int]] = {}
|
|
|
|
def update(self, boxes: Iterable[tuple[Box, float]], sequence: int) -> list[Detection]:
|
|
incoming = list(boxes)
|
|
available = set(self._tracks)
|
|
detections: list[Detection] = []
|
|
for box, score in incoming:
|
|
center = box.center
|
|
selected: str | None = None
|
|
selected_distance = self._max_distance
|
|
for track_id in available:
|
|
old_center = self._tracks[track_id][0].center
|
|
distance = ((center.x - old_center.x) ** 2 + (center.y - old_center.y) ** 2) ** 0.5
|
|
if distance < selected_distance:
|
|
selected = track_id
|
|
selected_distance = distance
|
|
if selected is None:
|
|
selected = f"P-{self._next_id:04d}"
|
|
self._next_id += 1
|
|
else:
|
|
available.remove(selected)
|
|
self._tracks[selected] = (box, sequence)
|
|
detections.append(Detection(selected, "person", box, score))
|
|
for track_id, (_, last_seen) in list(self._tracks.items()):
|
|
if sequence - last_seen > self._ttl_frames:
|
|
del self._tracks[track_id]
|
|
return detections
|