""" Detection Pipeline Worker Pulls frames from a StreamWorker queue and runs: 1. Vehicle detection (type, color, plate, person count) 2. Face detection (embedding, identity match) 3. Saves snapshots and events to DB """ import asyncio import cv2 import logging import os from datetime import datetime, timezone from pathlib import Path from uuid import UUID from sqlalchemy.ext.asyncio import AsyncSession from sqlalchemy import select, text from app.core.config import settings from app.core.database import AsyncSessionLocal from app.models.camera import Camera from app.models.vehicle import VehicleEvent, VehicleIdentity, VehicleType from app.models.person import FaceEvent, PersonIdentity from app.services.vehicle_detector import vehicle_detector from app.services.face_service import face_service from app.services.stream_manager import StreamWorker logger = logging.getLogger(__name__) class DetectionPipeline: """ Runs detection on frames for a single camera stream. """ def __init__(self, camera_id: UUID, camera_name: str, worker: StreamWorker): self.camera_id = camera_id self.camera_name = camera_name self.worker = worker self.snapshot_dir = Path(settings.STORAGE_PATH) / str(camera_id) self.snapshot_dir.mkdir(parents=True, exist_ok=True) self._running = False self._task = None async def start(self): self._running = True self._task = asyncio.create_task(self._loop()) logger.info(f"[{self.camera_name}] Pipeline started") async def stop(self): self._running = False if self._task: self._task.cancel() try: await self._task except asyncio.CancelledError: pass async def _loop(self): loop = asyncio.get_event_loop() while self._running: try: frame_data = await asyncio.wait_for( self.worker.frame_queue.get(), timeout=2.0 ) except asyncio.TimeoutError: continue except asyncio.CancelledError: raise frame = frame_data["frame"] captured_at = frame_data["captured_at"] try: await loop.run_in_executor( None, self._process_frame, frame, captured_at ) except Exception as e: logger.exception(f"[{self.camera_name}] Pipeline error: {e}") def _process_frame(self, frame, captured_at: datetime): """Synchronous processing — runs in thread pool.""" import asyncio loop = asyncio.new_event_loop() asyncio.set_event_loop(loop) try: loop.run_until_complete(self._async_process(frame, captured_at)) finally: loop.close() async def _async_process(self, frame, captured_at: datetime): """Run detections and save to DB.""" # --- Vehicle Detection --- vehicle_detections = vehicle_detector.detect(frame) # --- Face Detection --- face_detections = face_service.detect(frame) if not vehicle_detections and not face_detections: return # Save snapshot ts = captured_at.strftime("%Y%m%d_%H%M%S_%f") snapshot_filename = f"{ts}.jpg" snapshot_path = str(self.snapshot_dir / snapshot_filename) cv2.imwrite(snapshot_path, frame) async with AsyncSessionLocal() as db: vehicle_event_ids = [] # Save vehicle events for det in vehicle_detections: identity_id = await self._get_or_create_vehicle_identity( db, det.plate_number, det.vehicle_type, det.color ) # Save plate crop plate_snap = None if det.plate_crop is not None and det.plate_crop.size > 0: plate_snap = str(self.snapshot_dir / f"{ts}_plate.jpg") cv2.imwrite(plate_snap, det.plate_crop) event = VehicleEvent( camera_id=self.camera_id, identity_id=identity_id, plate_number=det.plate_number, plate_confidence=det.plate_confidence, vehicle_type=det.vehicle_type, color=det.color, person_count=det.person_count, detection_confidence=det.detection_confidence, vehicle_bbox=det.vehicle_bbox, plate_bbox=det.plate_bbox, snapshot_path=snapshot_path, plate_snapshot_path=plate_snap, captured_at=captured_at, ) db.add(event) await db.flush() vehicle_event_ids.append(event.id) # Save face events for i, det in enumerate(face_detections): identity_id, similarity = await self._match_or_create_person( db, det.embedding ) # Save face crop face_snap = None if det.face_crop is not None and det.face_crop.size > 0: face_snap = str(self.snapshot_dir / f"{ts}_face{i}.jpg") cv2.imwrite(face_snap, det.face_crop) embedding_list = face_service.embeddings_to_list(det.embedding) if det.embedding is not None else None event = FaceEvent( camera_id=self.camera_id, identity_id=identity_id, detection_confidence=det.detection_confidence, face_bbox=det.bbox, face_embedding=embedding_list, similarity_score=similarity, vehicle_event_id=vehicle_event_ids[0] if vehicle_event_ids else None, snapshot_path=snapshot_path, face_snapshot_path=face_snap, captured_at=captured_at, ) db.add(event) await db.commit() async def _get_or_create_vehicle_identity( self, db: AsyncSession, plate_number, vehicle_type, color ): """Find existing vehicle identity by plate or create new one.""" if not plate_number: return None result = await db.execute( select(VehicleIdentity).where(VehicleIdentity.plate_number == plate_number) ) identity = result.scalar_one_or_none() if identity: identity.last_seen_at = datetime.now(timezone.utc) identity.total_sightings += 1 if color and not identity.color: identity.color = color else: identity = VehicleIdentity( plate_number=plate_number, vehicle_type=vehicle_type or VehicleType.unknown, color=color, first_seen_at=datetime.now(timezone.utc), last_seen_at=datetime.now(timezone.utc), total_sightings=1, ) db.add(identity) await db.flush() return identity.id async def _match_or_create_person(self, db: AsyncSession, embedding): """ Find closest matching person identity using pgvector cosine similarity. If no match above threshold, create new identity. Returns (identity_id, similarity_score). """ if embedding is None: return None, 0.0 emb_list = face_service.embeddings_to_list(embedding) emb_str = f"[{','.join(str(x) for x in emb_list)}]" # pgvector cosine distance (1 - cosine_similarity) threshold = settings.FACE_SIMILARITY_THRESHOLD # distance threshold result = await db.execute( text(f""" SELECT id, 1 - (embedding <=> '{emb_str}'::vector) AS similarity FROM person_identities WHERE embedding IS NOT NULL AND (embedding <=> '{emb_str}'::vector) < :threshold ORDER BY embedding <=> '{emb_str}'::vector LIMIT 1 """), {"threshold": threshold} ) row = result.fetchone() if row: identity_id, similarity = row # Update last seen await db.execute( text("UPDATE person_identities SET last_seen_at=NOW(), total_sightings=total_sightings+1 WHERE id=:id"), {"id": identity_id} ) return identity_id, float(similarity) else: # New unknown person identity = PersonIdentity( first_seen_at=datetime.now(timezone.utc), last_seen_at=datetime.now(timezone.utc), total_sightings=1, embedding=emb_list, ) db.add(identity) await db.flush() return identity.id, 0.0