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
251 lines
8.8 KiB
Python
251 lines
8.8 KiB
Python
"""
|
|
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
|