itsrishu02/cctv-person-detection-api
0
1"""2ZEEX AI - Recognizer3Per-frame recognition orchestrator used by the FastAPI backend. Keeps a tiny4IoU tracker so we only re-detect every N frames but still update labels every5frame -> smooth UI without paying detection cost.6"""7from __future__ import annotations8 9import sqlite310import time11from dataclasses import dataclass12from pathlib import Path13from typing import Dict, List, Optional, Tuple14 15import cv216import numpy as np17 18from face_pipeline import (19 FaceDetector,20 FaceEmbedder,21 KnownFaces,22 MatchResult,23 PersonDetector,24 draw_label,25 load_known_faces,26 match_embedding,27 resize_keep_aspect,28)29 30 31GREEN = (0, 200, 0)32RED = (0, 0, 220)33ORANGE = (0, 140, 240) # Person detected but no face visible34YELLOW = (0, 200, 220)35 36 37@dataclass38class TrackedFace:39 box: Tuple[int, int, int, int]40 label: str41 color: Tuple[int, int, int]42 worker_id: Optional[str]43 score: float44 last_seen: int # frame index45 46 47@dataclass48class TrackedPerson:49 """A person detection that has NO face matched inside it (back-of-head etc)."""50 box: Tuple[int, int, int, int]51 score: float52 last_seen: int53 54 55@dataclass56class FrameStats:57 fps: float = 0.058 detections: int = 059 known_count: int = 060 unknown_count: int = 061 person_no_face_count: int = 062 total_people_count: int = 063 no_person_frame_count: int = 064 empty_zone_threshold: int = 365 empty_zone: bool = True66 zone_status: str = "EMPTY_ZONE"67 68 69class EventLogger:70 """Writes recognition events to a text log and SQLite, with per-worker cooldown."""71 72 def __init__(self, log_file: Optional[str], db_file: Optional[str],73 camera_id: str, zone: str, cooldown_seconds: float):74 self.log_file = Path(log_file) if log_file else None75 self.db_file = Path(db_file) if db_file else None76 self.camera_id = camera_id77 self.zone = zone78 self.cooldown = cooldown_seconds79 self._last_logged: Dict[str, float] = {}80 self._db: Optional[sqlite3.Connection] = None81 if self.log_file:82 self.log_file.parent.mkdir(parents=True, exist_ok=True)83 if self.db_file:84 self.db_file.parent.mkdir(parents=True, exist_ok=True)85 self._db = sqlite3.connect(str(self.db_file), check_same_thread=False)86 self._db.execute(87 """88 CREATE TABLE IF NOT EXISTS events (89 id INTEGER PRIMARY KEY AUTOINCREMENT,90 ts TEXT NOT NULL,91 worker_id TEXT,92 name TEXT,93 camera_id TEXT,94 zone TEXT,95 score REAL96 )97 """98 )99 self._db.commit()100 101 def log(self, match: MatchResult) -> bool:102 """Returns True if the event was actually written (cooldown not active)."""103 key = match.worker_id or "UNKNOWN"104 now = time.time()105 last = self._last_logged.get(key, 0.0)106 if now - last < self.cooldown:107 return False108 self._last_logged[key] = now109 ts = time.strftime("%Y-%m-%d %H:%M:%S", time.localtime(now))110 wid = match.worker_id or ""111 name = match.name112 line = (113 f"{ts} | worker_id={wid or '-':<10} | name={name:<25} | "114 f"camera={self.camera_id} | zone={self.zone} | score={match.score:.3f}"115 )116 if self.log_file:117 with open(self.log_file, "a", encoding="utf-8") as f:118 f.write(line + "\n")119 if self._db is not None:120 self._db.execute(121 "INSERT INTO events(ts, worker_id, name, camera_id, zone, score)"122 " VALUES (?, ?, ?, ?, ?, ?)",123 (ts, wid, name, self.camera_id, self.zone, float(match.score)),124 )125 self._db.commit()126 return True127 128 def close(self) -> None:129 if self._db is not None:130 self._db.close()131 self._db = None132 133 134class FrameRecognizer:135 """Wraps detector + embedder + matcher with a small IoU tracker."""136 137 def __init__(self, cfg: dict, base_dir: Path):138 paths = cfg["paths"]139 rec_cfg = cfg["recognition"]140 log_cfg = cfg["logging"]141 stream_cfg = cfg["stream"]142 person_cfg = cfg.get("person_detection", {}) or {}143 144 self.process_width = int(rec_cfg.get("process_width", 0))145 self.detect_every_n = max(1, int(rec_cfg.get("detect_every_n_frames", 1)))146 self.threshold = float(rec_cfg["cosine_threshold"])147 # Minimum face size (pixels, processed-frame coords) below which we148 # don't bother running the embedder - SFace embeddings on tiny149 # crops are noise. Far faces are still BOXED so the operator sees them.150 self.min_face_for_embed = int(rec_cfg.get("min_face_for_embed", 28))151 self.empty_zone_threshold_frames = max(152 1,153 int(rec_cfg.get("empty_zone_threshold_frames", 3)),154 )155 156 self.detector = FaceDetector(157 str((base_dir / paths["yolo_face_model"]).resolve()),158 conf=float(rec_cfg["yolo_conf"]),159 imgsz=int(rec_cfg.get("yolo_imgsz", 1280)),160 )161 self.embedder = FaceEmbedder(162 str((base_dir / paths["sface_model"]).resolve())163 )164 self.known: KnownFaces = load_known_faces(165 str((base_dir / paths["encodings_file"]).resolve())166 )167 168 # Optional person detector (catches back-of-head / occluded face)169 self.person_enabled = bool(person_cfg.get("enabled", False))170 self.person_detector: Optional[PersonDetector] = None171 self.person_model_path = str((172 base_dir / person_cfg.get("model", "models/yolov8n.pt")173 ).resolve())174 self.person_conf = float(person_cfg.get("conf", 0.5))175 self.person_imgsz = int(person_cfg.get("imgsz", 960))176 self.person_min_height_px = int(person_cfg.get("min_height_px", 70))177 self.person_min_aspect_ratio = float(person_cfg.get("min_aspect_ratio", 1.4))178 self.person_max_aspect_ratio = float(person_cfg.get("max_aspect_ratio", 4.5))179 self.person_max_area_frac = float(person_cfg.get("max_area_frac", 0.55))180 181 self.logger: Optional[EventLogger] = None182 if log_cfg.get("enable_file_log") or log_cfg.get("enable_sqlite"):183 self.logger = EventLogger(184 log_file=str((base_dir / paths["events_log"]).resolve())185 if log_cfg.get("enable_file_log") else None,186 db_file=str((base_dir / paths["events_db"]).resolve())187 if log_cfg.get("enable_sqlite") else None,188 camera_id=str(stream_cfg.get("camera_id", "CAM")),189 zone=str(stream_cfg.get("zone", "")),190 cooldown_seconds=float(log_cfg.get("cooldown_seconds", 30)),191 )192 193 # State194 self.frame_idx = 0195 self.tracked: List[TrackedFace] = []196 self.tracked_persons: List[TrackedPerson] = []197 self._no_person_frame_count = 0198 self._t_prev = time.time()199 self._fps = 0.0200 201 @property202 def n_known_workers(self) -> int:203 return len(self.known.workers)204 205 def update_camera_zone(self, camera_id: Optional[str] = None,206 zone: Optional[str] = None) -> None:207 """Allow the UI to change the labels going into events.log on the fly."""208 if self.logger is None:209 return210 if camera_id is not None:211 self.logger.camera_id = camera_id212 if zone is not None:213 self.logger.zone = zone214 215 def close(self) -> None:216 if self.logger is not None:217 self.logger.close()218 219 def _ensure_person_detector(self) -> PersonDetector:220 if self.person_detector is None:221 self.person_detector = PersonDetector(222 self.person_model_path,223 conf=self.person_conf,224 imgsz=self.person_imgsz,225 min_height_px=self.person_min_height_px,226 min_aspect_ratio=self.person_min_aspect_ratio,227 max_aspect_ratio=self.person_max_aspect_ratio,228 max_area_frac=self.person_max_area_frac,229 )230 return self.person_detector231 232 def _label_for(self, m: MatchResult) -> Tuple[str, Tuple[int, int, int]]:233 if m.worker_id is None:234 return "UNKNOWN", RED235 rec = m.record236 zone_part = f" | {rec.zone}" if rec and rec.zone else ""237 return f"[{m.worker_id}] {m.name}{zone_part} ({m.score:.2f})", GREEN238 239 @staticmethod240 def _face_inside_person(face_box: Tuple[int, int, int, int],241 person_box: Tuple[int, int, int, int]) -> bool:242 """Is the face's CENTER inside the person box, AND does the face243 sit in roughly the upper half of the person body?"""244 fx1, fy1, fx2, fy2 = face_box245 px1, py1, px2, py2 = person_box246 cx = (fx1 + fx2) / 2.0247 cy = (fy1 + fy2) / 2.0248 if not (px1 <= cx <= px2 and py1 <= cy <= py2):249 return False250 # face should be above the person box's vertical midpoint251 return cy <= (py1 + (py2 - py1) * 0.65)252 253 def process(self, frame_bgr: np.ndarray) -> Tuple[np.ndarray, FrameStats]:254 """Process a single BGR frame. Returns (annotated_frame, stats)."""255 self.frame_idx += 1256 proc, scale = resize_keep_aspect(frame_bgr, self.process_width)257 258 is_detect_frame = (self.frame_idx == 1259 or (self.frame_idx % self.detect_every_n == 0))260 261 if is_detect_frame:262 # ---- Face detection + recognition ----263 face_boxes = self.detector.detect(proc)264 new_tracked: List[TrackedFace] = []265 for (x1, y1, x2, y2, _conf) in face_boxes:266 bw = x2 - x1; bh = y2 - y1267 # Tiny faces: still draw a box so the operator sees them, but268 # don't try to ID them - SFace on a 15-px crop is just noise.269 if min(bw, bh) < self.min_face_for_embed:270 new_tracked.append(TrackedFace(271 box=(x1, y1, x2, y2),272 label="FACE (too far)",273 color=YELLOW,274 worker_id=None,275 score=0.0,276 last_seen=self.frame_idx,277 ))278 continue279 emb = self.embedder.embed(proc, (x1, y1, x2, y2))280 if emb is None:281 continue282 m = match_embedding(emb, self.known, self.threshold)283 label, color = self._label_for(m)284 new_tracked.append(TrackedFace(285 box=(x1, y1, x2, y2),286 label=label,287 color=color,288 worker_id=m.worker_id,289 score=m.score,290 last_seen=self.frame_idx,291 ))292 if self.logger is not None:293 self.logger.log(m)294 self.tracked = new_tracked295 296 # ---- Person detection (back-of-head fallback) ----297 new_persons: List[TrackedPerson] = []298 if self.person_enabled:299 person_boxes = self._ensure_person_detector().detect(proc)300 for (px1, py1, px2, py2, pconf) in person_boxes:301 has_face = any(302 self._face_inside_person(t.box, (px1, py1, px2, py2))303 for t in self.tracked304 )305 if has_face:306 continue # face label already covers this person307 new_persons.append(TrackedPerson(308 box=(px1, py1, px2, py2),309 score=pconf,310 last_seen=self.frame_idx,311 ))312 self.tracked_persons = new_persons313 # else: keep previous self.tracked / self.tracked_persons314 315 # FPS (EMA)316 now = time.time()317 dt = now - self._t_prev318 self._t_prev = now319 if dt > 0:320 inst = 1.0 / dt321 self._fps = inst if self._fps == 0 else (0.8 * self._fps + 0.2 * inst)322 323 # Draw on the original frame; upscale boxes from proc-coords if needed.324 out = frame_bgr.copy()325 inv = 1.0 / scale if scale != 0 else 1.0326 327 # 1) draw "person without face" boxes FIRST (so face labels go on top328 # if there's any overlap/border)329 for tp in self.tracked_persons:330 x1, y1, x2, y2 = tp.box331 if scale != 1.0:332 x1 = int(x1 * inv); y1 = int(y1 * inv)333 x2 = int(x2 * inv); y2 = int(y2 * inv)334 draw_label(out, (x1, y1, x2, y2),335 f"PERSON (no face) ({tp.score:.2f})", ORANGE)336 337 # 2) draw face boxes338 known_n = unknown_n = far_n = 0339 for t in self.tracked:340 x1, y1, x2, y2 = t.box341 if scale != 1.0:342 x1 = int(x1 * inv); y1 = int(y1 * inv)343 x2 = int(x2 * inv); y2 = int(y2 * inv)344 draw_label(out, (x1, y1, x2, y2), t.label, t.color)345 if t.color == YELLOW:346 far_n += 1347 elif t.worker_id is None:348 unknown_n += 1349 else:350 known_n += 1351 352 person_no_face_n = len(self.tracked_persons)353 total_people_n = known_n + unknown_n + far_n + person_no_face_n354 if total_people_n == 0:355 self._no_person_frame_count += 1356 else:357 self._no_person_frame_count = 0358 empty_zone = self._no_person_frame_count >= self.empty_zone_threshold_frames359 if empty_zone:360 zone_status = "EMPTY_ZONE"361 elif self._no_person_frame_count > 0:362 zone_status = "CHECKING_EMPTY_ZONE"363 else:364 zone_status = "OCCUPIED_ZONE"365 366 # HUD367 hud = (f"status:{zone_status} FPS:{self._fps:5.1f} faces:{len(self.tracked)} "368 f"(known:{known_n} unknown:{unknown_n} far:{far_n}) "369 f"persons-no-face:{person_no_face_n} "370 f"empty:{self._no_person_frame_count}/{self.empty_zone_threshold_frames} "371 f"thr:{self.threshold:.2f}")372 cv2.rectangle(out, (0, 0), (out.shape[1], 24), (32, 32, 32), -1)373 cv2.putText(out, hud, (8, 17),374 cv2.FONT_HERSHEY_SIMPLEX, 0.55, (255, 255, 255), 1, cv2.LINE_AA)375 376 if empty_zone:377 text = "EMPTY ZONE / NO PERSON"378 (tw, th), _bl = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 1.0, 2)379 pad_x = 18380 pad_y = 12381 x1 = max(0, (out.shape[1] - tw) // 2 - pad_x)382 y1 = max(30, (out.shape[0] - th) // 2 - pad_y)383 x2 = min(out.shape[1] - 1, x1 + tw + 2 * pad_x)384 y2 = min(out.shape[0] - 1, y1 + th + 2 * pad_y)385 cv2.rectangle(out, (x1, y1), (x2, y2), (22, 96, 130), -1)386 cv2.rectangle(out, (x1, y1), (x2, y2), (0, 190, 255), 2)387 cv2.putText(out, text, (x1 + pad_x, y2 - pad_y),388 cv2.FONT_HERSHEY_SIMPLEX, 1.0, (255, 255, 255), 2, cv2.LINE_AA)389 390 return out, FrameStats(391 fps=self._fps,392 detections=len(self.tracked),393 known_count=known_n,394 unknown_count=unknown_n,395 person_no_face_count=person_no_face_n,396 total_people_count=total_people_n,397 no_person_frame_count=self._no_person_frame_count,398 empty_zone_threshold=self.empty_zone_threshold_frames,399 empty_zone=empty_zone,400 zone_status=zone_status,401 )402 