7c43e4580d
- FastAPI backend with async SQLAlchemy - Camera RTSP management (add, start, stop) - Vehicle detection (YOLO + fast-alpr) - Type: car, motorcycle, truck, jeepney - Color detection (HSV) - License plate OCR - Motorcycle person count - Face detection + InsightFace ArcFace embedding - pgvector identity grouping (auto-cluster same face) - Vehicle + person movement trail APIs - Docker Compose with pgvector/pg16 - Models: Camera, VehicleIdentity, VehicleEvent, PersonIdentity, FaceEvent
154 lines
4.8 KiB
Python
154 lines
4.8 KiB
Python
"""
|
|
StreamManager — Manages RTSP stream connections and dispatches frames to workers.
|
|
Each camera runs in its own asyncio task with a frame queue.
|
|
"""
|
|
import asyncio
|
|
import cv2
|
|
import logging
|
|
from typing import Dict, Optional
|
|
from uuid import UUID
|
|
from datetime import datetime
|
|
|
|
from app.core.config import settings
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class StreamWorker:
|
|
"""
|
|
Handles one RTSP camera stream.
|
|
Reads frames and puts them into a queue for the AI pipeline to consume.
|
|
"""
|
|
|
|
def __init__(self, camera_id: UUID, camera_name: str, rtsp_url: str):
|
|
self.camera_id = camera_id
|
|
self.camera_name = camera_name
|
|
self.rtsp_url = rtsp_url
|
|
self.frame_queue: asyncio.Queue = asyncio.Queue(maxsize=10)
|
|
self.is_running = False
|
|
self._task: Optional[asyncio.Task] = None
|
|
self.frame_count = 0
|
|
self.last_frame_at: Optional[datetime] = None
|
|
self.error: Optional[str] = None
|
|
|
|
async def start(self):
|
|
self.is_running = True
|
|
self._task = asyncio.create_task(self._read_loop())
|
|
logger.info(f"[{self.camera_name}] Stream started")
|
|
|
|
async def stop(self):
|
|
self.is_running = False
|
|
if self._task:
|
|
self._task.cancel()
|
|
try:
|
|
await self._task
|
|
except asyncio.CancelledError:
|
|
pass
|
|
logger.info(f"[{self.camera_name}] Stream stopped")
|
|
|
|
async def _read_loop(self):
|
|
"""Read frames from RTSP in a thread pool (OpenCV is blocking)."""
|
|
loop = asyncio.get_event_loop()
|
|
|
|
while self.is_running:
|
|
try:
|
|
cap = await loop.run_in_executor(
|
|
None,
|
|
lambda: cv2.VideoCapture(self.rtsp_url)
|
|
)
|
|
|
|
if not cap.isOpened():
|
|
self.error = "Cannot open RTSP stream"
|
|
logger.error(f"[{self.camera_name}] {self.error}")
|
|
await asyncio.sleep(5)
|
|
continue
|
|
|
|
self.error = None
|
|
frame_idx = 0
|
|
|
|
while self.is_running:
|
|
ret, frame = await loop.run_in_executor(None, cap.read)
|
|
|
|
if not ret:
|
|
logger.warning(f"[{self.camera_name}] Frame read failed, reconnecting...")
|
|
break
|
|
|
|
frame_idx += 1
|
|
self.frame_count += 1
|
|
self.last_frame_at = datetime.utcnow()
|
|
|
|
# Skip frames for performance
|
|
if frame_idx % settings.PROCESS_EVERY_N_FRAMES != 0:
|
|
continue
|
|
|
|
# Non-blocking put — drop frame if queue is full
|
|
try:
|
|
self.frame_queue.put_nowait({
|
|
"frame": frame,
|
|
"frame_idx": frame_idx,
|
|
"captured_at": self.last_frame_at,
|
|
})
|
|
except asyncio.QueueFull:
|
|
pass # AI pipeline is slow — drop frame
|
|
|
|
cap.release()
|
|
|
|
except asyncio.CancelledError:
|
|
raise
|
|
except Exception as e:
|
|
self.error = str(e)
|
|
logger.exception(f"[{self.camera_name}] Stream error: {e}")
|
|
await asyncio.sleep(5)
|
|
|
|
|
|
class StreamManager:
|
|
"""
|
|
Singleton that manages all active camera streams.
|
|
"""
|
|
_instance = None
|
|
|
|
def __init__(self):
|
|
self._workers: Dict[str, StreamWorker] = {}
|
|
|
|
@classmethod
|
|
def get(cls) -> "StreamManager":
|
|
if cls._instance is None:
|
|
cls._instance = StreamManager()
|
|
return cls._instance
|
|
|
|
async def start_stream(self, camera_id: UUID, camera_name: str, rtsp_url: str):
|
|
key = str(camera_id)
|
|
if key in self._workers:
|
|
await self.stop_stream(camera_id)
|
|
|
|
worker = StreamWorker(camera_id, camera_name, rtsp_url)
|
|
self._workers[key] = worker
|
|
await worker.start()
|
|
return worker
|
|
|
|
async def stop_stream(self, camera_id: UUID):
|
|
key = str(camera_id)
|
|
if key in self._workers:
|
|
await self._workers[key].stop()
|
|
del self._workers[key]
|
|
|
|
def get_worker(self, camera_id: UUID) -> Optional[StreamWorker]:
|
|
return self._workers.get(str(camera_id))
|
|
|
|
def get_all_status(self) -> list:
|
|
return [
|
|
{
|
|
"camera_id": str(k),
|
|
"is_running": w.is_running,
|
|
"frame_count": w.frame_count,
|
|
"last_frame_at": w.last_frame_at,
|
|
"error": w.error,
|
|
}
|
|
for k, w in self._workers.items()
|
|
]
|
|
|
|
async def stop_all(self):
|
|
for worker in list(self._workers.values()):
|
|
await worker.stop()
|
|
self._workers.clear()
|