CoolFace
Apppublic

itsrishu02/cctv-person-detection-api

sourceHugging Faceupdated 5mo agoView on Hugging Face
0likes
recognizer.py402 linesDownload Raw Back to backend
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