faces.py

"""faces.py — Person tracking + face cropping for Marquee.

Performance fixes vs previous version:
  - Face detection runs ONCE per sampled frame (not once per person box)
  - Track lookup uses dict instead of linear scan — O(1) vs O(n)
  - Sightings stored only in JSONL, not accumulated in memory per track
  - Logging added throughout for visibility
"""

import json
import logging
import time
from pathlib import Path

import cv2
import numpy as np
from PIL import Image

from ov_models import OVModel

log = logging.getLogger(__name__)

PERSON_CONF   = 0.50
FACE_CONF     = 0.55
SAMPLE_FPS    = 2.0
IOU_THRESHOLD = 0.35
MAX_MISS      = 4


_person_det = None
_face_det   = None


def _load():
    global _person_det, _face_det
    if _person_det is None:
        log.info("[faces] Loading OpenVINO models (person + face detection)…")
        t0 = time.time()
        _person_det = OVModel("person-detection-retail-0013")
        _face_det   = OVModel("face-detection-retail-0004")
        log.info(f"[faces] Models loaded in {time.time()-t0:.1f}s")
    return _person_det, _face_det


def _detect_persons(frame_bgr, det) -> list[list[int]]:
    h, w = frame_bgr.shape[:2]
    out  = det.infer(frame_bgr).reshape(-1, 7)
    boxes = []
    for _, _, conf, x1, y1, x2, y2 in out:
        if conf < PERSON_CONF:
            continue
        boxes.append([max(0,int(x1*w)), max(0,int(y1*h)),
                      min(w,int(x2*w)), min(h,int(y2*h))])
    return boxes


def _detect_faces_once(frame_bgr, fdet) -> list[list[int]]:
    """Run face detection ONCE per frame — reuse results for all person boxes."""
    h, w = frame_bgr.shape[:2]
    out  = fdet.infer(frame_bgr).reshape(-1, 7)
    boxes = []
    for _, _, conf, x1, y1, x2, y2 in out:
        if conf < FACE_CONF:
            continue
        boxes.append([max(0,int(x1*w)), max(0,int(y1*h)),
                      min(w,int(x2*w)), min(h,int(y2*h))])
    return boxes


def _best_face_for_person(person_box, all_faces) -> list[int] | None:
    """Match a person box to the face with highest IoU from the pre-detected list."""
    px1, py1, px2, py2 = person_box
    best_iou, best_box = 0.0, None
    for fx1, fy1, fx2, fy2 in all_faces:
        ix1=max(px1,fx1); iy1=max(py1,fy1)
        ix2=min(px2,fx2); iy2=min(py2,fy2)
        inter = max(0,ix2-ix1)*max(0,iy2-iy1)
        if inter == 0:
            continue
        union = (px2-px1)*(py2-py1) + (fx2-fx1)*(fy2-fy1) - inter
        iou   = inter / (union + 1e-6)
        if iou > best_iou:
            best_iou, best_box = iou, [fx1,fy1,fx2,fy2]
    return best_box


def _iou(a, b) -> float:
    ax1,ay1,ax2,ay2 = a
    bx1,by1,bx2,by2 = b
    ix1=max(ax1,bx1); iy1=max(ay1,by1)
    ix2=min(ax2,bx2); iy2=min(ay2,by2)
    inter = max(0,ix2-ix1)*max(0,iy2-iy1)
    if inter == 0: return 0.0
    return inter/((ax2-ax1)*(ay2-ay1)+(bx2-bx1)*(by2-by1)-inter+1e-6)


class _Track:
    _next_id = 0

    def __init__(self, box):
        self.id         = _Track._next_id; _Track._next_id += 1
        self.box        = box
        self.miss       = 0
        self.best_crop: Image.Image | None = None
        self.best_area  = 0
        self.count      = 0

    def update(self, box, crop: Image.Image | None):
        self.box   = box
        self.miss  = 0
        self.count += 1
        if crop is not None:
            area = crop.size[0] * crop.size[1]
            if area > self.best_area:
                self.best_crop = crop
                self.best_area = area


class _IoUTracker:
    def __init__(self):
        self.tracks: list[_Track] = []
        self._id_map: dict[int, _Track] = {}  # O(1) lookup by track id

    def step(self, boxes, crops) -> list[tuple[int, list[int], _Track]]:
        matched_t = set()
        matched_d = set()
        pairs = []

        for di, dbox in enumerate(boxes):
            best_iou, best_ti = 0.0, -1
            for ti, t in enumerate(self.tracks):
                iou = _iou(t.box, dbox)
                if iou > best_iou:
                    best_iou, best_ti = iou, ti
            if best_iou >= IOU_THRESHOLD and best_ti not in matched_t:
                pairs.append((best_ti, di))
                matched_t.add(best_ti); matched_d.add(di)

        for ti, di in pairs:
            self.tracks[ti].update(boxes[di], crops[di])

        for ti, t in enumerate(self.tracks):
            if ti not in matched_t:
                t.miss += 1

        for di, (dbox, crop) in enumerate(zip(boxes, crops)):
            if di not in matched_d:
                t = _Track(dbox); t.update(dbox, crop)
                self.tracks.append(t)
                self._id_map[t.id] = t

        self.tracks = [t for t in self.tracks if t.miss <= MAX_MISS]
        # keep id_map in sync
        alive = {t.id for t in self.tracks}
        self._id_map = {k: v for k, v in self._id_map.items() if k in alive}

        return [(t.id, t.box, t) for t in self.tracks if t.miss == 0]

    def get_by_id(self, tid: int) -> "_Track | None":
        return self._id_map.get(tid)


def scan_video(video_path: str, max_seconds: float = 600.0):
    """Person tracking scan with progress logging.

    Returns:
      tracks:     [{id, crop (PIL), sightings: [{t, bbox}]}]  — strong tracks only
      motion:     {frame_idx: float}
      meta:       {"fps", "n_frames"}
      jsonl_path: str
    """
    t_start = time.time()
    log.info(f"[scan] Starting scan: {video_path}")

    pdet, fdet = _load()
    _Track._next_id = 0

    cap    = cv2.VideoCapture(video_path)
    fps    = cap.get(cv2.CAP_PROP_FPS) or 30.0
    total  = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) or 0
    stride = max(1, int(round(fps / SAMPLE_FPS)))

    log.info(f"[scan] Video: {total} frames @ {fps:.1f}fps, "
             f"sampling every {stride} frames ({SAMPLE_FPS}fps)")

    tracker     = _IoUTracker()
    all_tracks: dict[int, dict] = {}
    motion: dict[int, float]    = {}
    prev_small  = None
    idx         = 0
    sampled     = 0
    persons_seen = 0

    while True:
        ok, frame = cap.read()
        if not ok or idx / fps > max_seconds:
            break

        if idx % stride == 0:
            t_sec = round(idx / fps, 3)
            sampled += 1

            # motion
            small = cv2.cvtColor(cv2.resize(frame,(160,90)), cv2.COLOR_BGR2GRAY)
            if prev_small is not None:
                motion[idx] = float(np.mean(cv2.absdiff(small, prev_small)))
            prev_small = small

            # detect persons + faces (ONE face inference per frame, not per person)
            p_boxes = _detect_persons(frame, pdet)
            f_boxes = _detect_faces_once(frame, fdet) if p_boxes else []
            persons_seen += len(p_boxes)

            # build face crops per person
            crops = []
            for pb in p_boxes:
                fb = _best_face_for_person(pb, f_boxes)
                if fb is not None:
                    fx1,fy1,fx2,fy2 = fb
                    face_img = Image.fromarray(
                        cv2.cvtColor(frame[fy1:fy2,fx1:fx2], cv2.COLOR_BGR2RGB))
                else:
                    px1,py1,px2,py2 = pb
                    fh = max(1,(py2-py1)//3)
                    face_img = Image.fromarray(
                        cv2.cvtColor(frame[py1:py1+fh,px1:px2], cv2.COLOR_BGR2RGB))
                crops.append(face_img)

            # track
            active = tracker.step(p_boxes, crops)
            for tid, bbox, trk in active:
                if tid not in all_tracks:
                    all_tracks[tid] = {"id": tid, "crop": None,
                                       "sightings": [], "count": 0}
                entry = all_tracks[tid]
                entry["count"] += 1
                entry["sightings"].append({"t": t_sec, "bbox": bbox})
                if trk.best_crop is not None:
                    area = trk.best_crop.size[0] * trk.best_crop.size[1]
                    if entry["crop"] is None or area > entry.get("_area", 0):
                        entry["crop"]  = trk.best_crop
                        entry["_area"] = area

            # progress log every ~5 seconds of video
            if sampled % max(1, int(5 * SAMPLE_FPS)) == 0:
                pct = int(100 * idx / max(total, 1))
                log.info(f"[scan] {pct}% — t={t_sec:.1f}s, "
                         f"active tracks={len(tracker.tracks)}, "
                         f"persons this frame={len(p_boxes)}")

        idx += 1

    cap.release()
    elapsed = time.time() - t_start

    strong = {tid: v for tid, v in all_tracks.items() if v["count"] >= 2}
    if not strong:
        strong = all_tracks

    log.info(f"[scan] Done in {elapsed:.1f}s — "
             f"sampled {sampled} frames, "
             f"total tracks={len(all_tracks)}, "
             f"strong tracks={len(strong)}, "
             f"total person detections={persons_seen}")

    # write JSONL
    jsonl_path = str(Path(video_path).with_suffix(".tracks.jsonl"))
    rows_written = 0
    with open(jsonl_path, "w") as f:
        for tid, entry in strong.items():
            for sight in entry["sightings"]:
                f.write(json.dumps({
                    "track_id": tid,
                    "t": sight["t"],
                    "bbox": sight["bbox"],
                    "name": "",
                }) + "\n")
                rows_written += 1
    log.info(f"[scan] JSONL written: {jsonl_path} ({rows_written} rows)")

    tracks_list = [
        {"id": v["id"], "crop": v["crop"], "sightings": v["sightings"]}
        for v in strong.values()
    ]
    return tracks_list, motion, {"fps": fps, "n_frames": idx}, jsonl_path


def write_names_to_jsonl(jsonl_path: str, id_to_name: dict[int, str]):
    log.info(f"[faces] Writing names to JSONL: {id_to_name}")
    rows = []
    with open(jsonl_path) as f:
        for line in f:
            row = json.loads(line)
            tid = row["track_id"]
            if tid in id_to_name:
                row["name"] = id_to_name[tid]
            rows.append(row)
    with open(jsonl_path, "w") as f:
        for row in rows:
            f.write(json.dumps(row) + "\n")
    log.info(f"[faces] Names written for {len(id_to_name)} tracks")


def annotate_event_frame(frame_bgr, jsonl_path: str, t_sec: float,
                         tol: float = 1.5) -> Image.Image:
    img = frame_bgr.copy()
    try:
        entries = []
        with open(jsonl_path) as f:
            for line in f:
                row = json.loads(line)
                if abs(row["t"] - t_sec) <= tol and row.get("name","").strip():
                    entries.append(row)
        if entries:
            log.debug(f"[annotate] t={t_sec:.2f}s — burning {len(entries)} names")
    except (FileNotFoundError, json.JSONDecodeError) as e:
        log.warning(f"[annotate] Could not read JSONL: {e}")
        return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))

    for row in entries:
        name = row["name"].strip()
        if not name: continue
        x1,y1,x2,y2 = row["bbox"]
        scale = max(0.5,(x2-x1)/200.0)
        thick = max(1,int(scale*2))
        (tw,th),_ = cv2.getTextSize(name, cv2.FONT_HERSHEY_DUPLEX, scale, thick)
        tx=x1; ty=max(th+6,y1-8)
        cv2.rectangle(img,(tx-4,ty-th-6),(tx+tw+4,ty+4),(0,0,0),-1)
        cv2.putText(img,name,(tx,ty),cv2.FONT_HERSHEY_DUPLEX,
                    scale,(0,255,255),thick,cv2.LINE_AA)

    return Image.fromarray(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))


def frame_at(video_path: str, frame_idx: int):
    cap = cv2.VideoCapture(video_path)
    cap.set(cv2.CAP_PROP_POS_FRAMES, frame_idx)
    ok, frame = cap.read()
    cap.release()
    if not ok:
        log.warning(f"[faces] frame_at({frame_idx}) failed")
    return frame if ok else None