Behavision/behavision/tracking.py

"""IoU-based multi-face tracker.

Purpose: turn per-frame detections into per-person *tracks* so identity is
decided once per visit, not once per frame (the old backend registered a
new user for every frame). Greedy IoU association is deliberate: faces move
slowly relative to frame rate, and determinism beats a heavier Kalman/
ByteTrack stack for this workload.
"""
from __future__ import annotations

import itertools
import time
from dataclasses import dataclass, field
from typing import Optional

import numpy as np

from .detection import Detection
from .geometry import iou


@dataclass
class Track:
    id: int
    box: tuple
    kps: np.ndarray
    score: float
    quality: float = 0.0
    best_quality: float = 0.0
    hits: int = 1
    misses: int = 0
    created_at: float = field(default_factory=time.time)
    updated_at: float = field(default_factory=time.time)
    # identity resolution state
    state: str = "pending"  # pending | resolved | ambiguous | gave_up
    # Embeddings are accumulated over multiple frames and averaged before
    # any identity decision: single-frame embeddings under extreme pose /
    # motion blur are unstable, the mean is not.
    emb_sum: Optional[np.ndarray] = None
    emb_count: int = 0
    id_attempts: int = 0
    # Set when THIS track minted the identity, so a terminal tally can tell
    # a first-time visitor from a returning one without re-querying the store.
    is_new: bool = False
    # resolve() refused to enroll this face (quality below the gate). Counted
    # rather than ignored: a mis-set gate and an empty room used to look the
    # same from outside.
    quality_skips: int = 0
    last_attempt_ts: float = 0.0
    last_reinforce_ts: float = 0.0
    reinforcements: int = 0
    identity_id: Optional[int] = None
    label: Optional[str] = None
    similarity: float = 0.0
    attributes: dict = field(default_factory=dict)
    # The best-quality face crop seen on this track, kept only when
    # app.store_faces is on. One small array per live track, replaced rather
    # than accumulated; None when images are off, which is the default.
    best_face: Optional[np.ndarray] = None
    best_face_quality: float = 0.0
    attr_samples: list = field(default_factory=list)


class IouTracker:
    def __init__(self, iou_threshold: float = 0.3, max_misses: int = 15):
        self.iou_threshold = iou_threshold
        self.max_misses = max_misses
        self.tracks: list[Track] = []
        self._ids = itertools.count(1)

    def update(self, detections: "list[Detection]", now: "float | None" = None
               ) -> "tuple[list[Track], list[Track]]":
        """Associate detections to tracks. Returns (active, ended)."""
        now = now or time.time()

        # Greedy matching on IoU, best pairs first.
        pairs = []
        for ti, track in enumerate(self.tracks):
            for di, det in enumerate(detections):
                overlap = iou(track.box, det.box)
                if overlap >= self.iou_threshold:
                    pairs.append((overlap, ti, di))
        pairs.sort(reverse=True)

        matched_tracks: set[int] = set()
        matched_dets: set[int] = set()
        for overlap, ti, di in pairs:
            if ti in matched_tracks or di in matched_dets:
                continue
            matched_tracks.add(ti)
            matched_dets.add(di)
            track, det = self.tracks[ti], detections[di]
            track.box = det.box
            track.kps = det.kps
            track.score = det.score
            track.quality = det.quality
            track.best_quality = max(track.best_quality, det.quality)
            track.hits += 1
            track.misses = 0
            track.updated_at = now

        new_tracks = [
            Track(id=next(self._ids), box=det.box, kps=det.kps,
                  score=det.score, quality=det.quality,
                  best_quality=det.quality, created_at=now, updated_at=now)
            for di, det in enumerate(detections) if di not in matched_dets
        ]

        ended: list[Track] = []
        alive: list[Track] = []
        for ti, track in enumerate(self.tracks):
            if ti not in matched_tracks:
                track.misses += 1
            if track.misses > self.max_misses:
                ended.append(track)
            else:
                alive.append(track)
        self.tracks = alive + new_tracks
        return self.tracks, ended