scene-actor-extraction/experiments/xsource/landmark_voting.py

#!/usr/bin/env python3
"""Landmark voting: average SCRFD's overlapping detections instead of discarding them.

SCRFD predicts a face from many anchors; NMS keeps the single highest-scoring
box and throws the rest away. Each discarded box carries its own 5-landmark
estimate of the SAME face, so the survivors are one sample from a distribution
we could be averaging over.

    baseline  conf 0.50, nms 0.40 — the shipped settings, one box per face
    voted     conf 0.30, nms 0.90 — duplicates survive, then grouped by IoU and
              the 5 landmarks averaged, weighted by detection confidence

Why this is worth trying when the mesh failed: the mesh moved the landmarks off
the definition ArcFace was trained on (a lip-ring centroid is not an annotated
mouth corner), and the embedder punished it. A confidence-weighted mean of
SCRFD's OWN landmark predictions is the same kind of point, just with less
variance — it should stay on-distribution while being steadier.

Scored on cross-clip identification through the production sigmoid, which is
the thing that actually broke. Raw similarity shown only to locate the
threshold; it decides nothing.

    LD_PRELOAD=/usr/lib/libcudnn_cnn.so.9 python3 landmark_voting.py
"""
import sys, glob, json, os

sys.path.insert(0, "/home/dtourolle/Development/Jray-project/scene-actor-extraction/build-ort")
import sae_embed          # before cv2 — see alignment_compare.py

import numpy as np
import cv2

ROOT = "/home/dtourolle/Development/Jray-project/scene-actor-extraction/"
M = ROOT + "models/"
CLIPS = ["5157344", "5157339"]
PROB_THRESHOLD = 0.754
GROUP_IOU = 0.55          # detections overlapping this much are the same face
MATCH_IOU = 0.35          # tie a detection to the hand-labelled face

base_eng = sae_embed.FaceEmbedder(detector_model=M + "scrfd_500m_bnkps.onnx",
                                  arcface_model=M + "LVFace-B_Glint360K.onnx",
                                  conf=0.5, nms=0.4, max_side=0)
# Same models, looser suppression: keep the duplicates NMS would have removed.
vote_eng = sae_embed.FaceEmbedder(detector_model=M + "scrfd_500m_bnkps.onnx",
                                  arcface_model=M + "LVFace-B_Glint360K.onnx",
                                  conf=0.3, nms=0.9, max_side=0)
cal = sae_embed.gallery_calibration(ROOT + "gallery_lvface.h5")


def iou(a, b):
    ax, ay, aw, ah = a; bx, by, bw, bh = b
    x0, y0 = max(ax, bx), max(ay, by)
    x1, y1 = min(ax + aw, bx + bw), min(ay + ah, by + bh)
    if x1 <= x0 or y1 <= y0:
        return 0.0
    i = (x1 - x0) * (y1 - y0)
    return i / (aw * ah + bw * bh - i)


def vote(dets):
    """Group overlapping detections, return (bbox, landmarks, conf, n_votes)."""
    items = sorted(dets, key=lambda d: -d.confidence)
    used, out = [False] * len(items), []
    for i, d in enumerate(items):
        if used[i]:
            continue
        grp = [d]
        used[i] = True
        for j in range(i + 1, len(items)):
            if not used[j] and iou(list(d.bbox), list(items[j].bbox)) >= GROUP_IOU:
                used[j] = True
                grp.append(items[j])
        w = np.array([g.confidence for g in grp], dtype=np.float32)
        w = w / w.sum()
        lms = np.stack([np.array(g.landmarks, dtype=np.float32).reshape(5, 2) for g in grp])
        bxs = np.stack([np.array(list(g.bbox), dtype=np.float32) for g in grp])
        out.append((( w[:, None] * bxs).sum(0), (w[:, None, None] * lms).sum(0),
                    float(grp[0].confidence), len(grp)))
    return out


def collect(clip):
    lab = {os.path.basename(p): os.path.basename(os.path.dirname(p))
           for p in glob.glob(f"labelling/{clip}/*/*.jpg")
           if os.path.basename(os.path.dirname(p)) not in ("discard", "unsorted")}
    man = {m["file"]: m for m in json.load(open(f"labelling/{clip}/manifest.json"))}
    rows, votes = [], []
    for frame in sorted({man[f]["frame"] for f in lab}):
        img = cv2.imread(f"frames/d{clip}_{frame}.png")
        base = base_eng.detect(img)
        voted = vote(vote_eng.detect(img))
        for fname, person in lab.items():
            m = man[fname]
            if m["frame"] != frame or m["idx"] >= len(base):
                continue
            d = base[m["idx"]]
            lm5 = np.array(d.landmarks, dtype=np.float32).reshape(5, 2)
            c_b = sae_embed.align_face(img, lm5)

            # the voted group covering the same face
            best, best_v = None, 0.0
            for bbox, lms, conf, n in voted:
                v = iou(list(bbox), list(d.bbox))
                if v > best_v:
                    best_v, best = v, (lms, n)
            c_v = None
            if best and best_v >= MATCH_IOU:
                c_v = sae_embed.align_face(img, best[0].astype(np.float32))
                votes.append(best[1])
            rec = {"person": person}
            rec["base"] = np.asarray(base_eng.embed_crop(c_b), np.float32) if c_b is not None else None
            rec["voted"] = np.asarray(base_eng.embed_crop(c_v), np.float32) if c_v is not None else None
            rows.append(rec)
    return rows, votes


data, allv = {}, []
for c in CLIPS:
    data[c], v = collect(c)
    allv += v
    print(f"[{c}] {len(data[c])} crops", file=sys.stderr)
print(f"[voting] group size: median {np.median(allv):.0f}, "
      f"mean {np.mean(allv):.1f}, max {max(allv)} detections averaged per face",
      file=sys.stderr)

GAL, PRB = "5157344", "5157339"
print(f"\ngallery {GAL} -> probe {PRB}, P>{PROB_THRESHOLD}\n")
print(f"{'align':>8}{'person':>8}{'n_gal':>7}{'n_prb':>7}"
      f"{'within-clip':>13}{'cross-clip':>12}{'hit rate':>10}")
summary = {}
for key in ("base", "voted"):
    gal, prb = {}, {}
    for r in data[GAL]:
        if r[key] is not None:
            gal.setdefault(r["person"], []).append(r[key])
    for r in data[PRB]:
        if r[key] is not None:
            prb.setdefault(r["person"], []).append(r[key])
    gal = {p: np.stack(v) for p, v in gal.items()}
    prb = {p: np.stack(v) for p, v in prb.items()}
    hits = tot = 0
    for p in sorted(set(gal) & set(prb)):
        pp = prb[p] @ prb[p].T
        np.fill_diagonal(pp, -1)
        within = float(np.median(pp.max(axis=1))) if len(pp) > 1 else float("nan")
        cross = float(np.median((gal[p] @ prb[p].T).max(axis=0)))
        h = 0
        for e in prb[p]:
            bp, bn = 0.0, None
            for q in gal:
                v = cal.probability(float((gal[q] @ e).max()))
                if v > bp:
                    bp, bn = v, q
            if bp > PROB_THRESHOLD and bn == p:
                h += 1
        hits += h; tot += len(prb[p])
        print(f"{key:>8}{p:>8}{len(gal[p]):>7}{len(prb[p]):>7}"
              f"{cal.probability(within):>6.3f}/{within:<6.3f}"
              f"{cal.probability(cross):>6.3f}/{cross:<5.3f}{100*h/len(prb[p]):>9.0f}%")
    summary[key] = (hits, tot)
    print(f"{key:>8}{'ALL':>8}{'':>14}{'':>25}{100*hits/max(tot,1):>9.0f}%\n")

hb, tb = summary["base"]; hv, tv = summary["voted"]
print(f"voting vs baseline: {100*hv/max(tv,1) - 100*hb/max(tb,1):+.1f} points "
      f"of cross-clip TPI ({hb}/{tb} -> {hv}/{tv})")