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

#!/usr/bin/env python3
"""Propose person labels for one clip using another clip's hand-sorted labels.

Reads the clip you have already sorted (REF_CLIP) as ground truth, then proposes
a person for every crop in the other clip (TARGET_CLIP) and writes them into
matching folders for you to correct.

    python3 propose_labels.py                  # propose, write folders + sheets
    python3 propose_labels.py --dry-run        # report only, move nothing

Output:
    labelling/<target>/unsorted/A|B|C|D/       proposed, same names as the ref clip
    labelling/<target>/unsorted/               left in place when no person is
                                               confident enough to name
    labelling/review_<person>.jpg              contact sheet spanning BOTH clips:
                                               confirmed crops first, then
                                               proposed ones with their P

Correcting it: open a review sheet. Every face on it should be one person. The
lower block is the proposal — move any intruder to the right folder, or back to
unsorted/. The folder a file sits in is the ground truth; nothing downstream
reads the proposed name or its probability.

The proposal is a labelling aid, never the label. Scoring the sweep against
embedding-derived labels would be circular: it keeps the faces the embedder
already gets right and drops the hard ones the sweep exists to find. Your
correction is what breaks that loop, which is why the proposal is deliberately
conservative and leaves anything doubtful unnamed.

Assignment is on the calibrated probability, per-actor best-of-N, exactly as
identity_matcher_node does — never a bare cosine (AR-024). The calibration is
fitted on your labelled reference crops, which is what calibrate_gallery is for.
"""
import sys, glob, json, os, shutil
import numpy as np
import cv2

sys.path.insert(0, "/home/dtourolle/Development/Jray-project/scene-actor-extraction/build-ort")
import sae_embed

ROOT = "/home/dtourolle/Development/Jray-project/scene-actor-extraction/"
M = ROOT + "models/"
# The embedder and the gallery whose calibration scores it MUST be the same
# model: a Platt fit is specific to one embedding space, so LVFace probabilities
# read through an ArcFace fit are meaningless.
EMBEDDER = M + "LVFace-B_Glint360K.onnx"
GALLERY  = ROOT + "gallery_lvface.h5"      # 291 actors, cached fit
REF_CLIP, TARGET_CLIP = "5157344", "5157339"
ASSIGN_P  = 0.90      # propose a name only when this confident
SHEET_COLS = 8
THUMB = 150
DRY = "--dry-run" in sys.argv

eng = sae_embed.FaceEmbedder(detector_model=M + "scrfd_500m_bnkps.onnx",
                             arcface_model=EMBEDDER,
                             conf=0.5, nms=0.4, max_side=0)


def embed_manifest(clip):
    """Re-derive each dumped crop's embedding from its source frame, cached.

    The dumped .jpg is a context thumbnail for human eyes; the embedding must
    come from the aligned crop the pipeline would actually produce, so the
    frame is re-detected and the manifest's idx picks the same face.

    Detecting 24 4K frames per clip costs far more than the rest of this script
    put together, and the result only changes when the manifest does — so it is
    cached and keyed on the manifest's mtime. Delete cache/ to force a redo.
    """
    man_path = f"labelling/{clip}/manifest.json"
    cache_path = f"cache/emb_{clip}.npz"
    os.makedirs("cache", exist_ok=True)
    if os.path.exists(cache_path) and \
       os.path.getmtime(cache_path) >= os.path.getmtime(man_path):
        z = np.load(cache_path, allow_pickle=True)
        print(f"[cache] {clip}: {len(z['meta'])} embeddings reused", file=sys.stderr)
        return [{**m, "emb": e} for m, e in zip(z["meta"], z["emb"])]

    man = json.load(open(man_path))
    by_frame = {}
    for m in man:
        by_frame.setdefault(m["frame"], []).append(m)
    out = []
    for frame, ms in sorted(by_frame.items()):
        img = cv2.imread(f"frames/d{clip}_{frame}.png")
        if img is None:
            sys.exit(f"missing frames/d{clip}_{frame}.png — extract with\n"
                     f"  ffmpeg -i clips/{clip}.mp4 -vf fps=2 -frames:v 24 "
                     f"frames/d{clip}_%03d.png")
        dets = eng.detect(img)
        for m in ms:
            if m["idx"] >= len(dets):
                continue
            d = dets[m["idx"]]
            lm = np.array(d.landmarks, dtype=np.float32).reshape(5, 2)
            crop = sae_embed.align_face(img, lm)
            if crop is None:
                continue
            out.append({**m, "emb": np.asarray(eng.embed_crop(crop), dtype=np.float32)})

    np.savez(cache_path,
             meta=np.array([{k: v for k, v in o.items() if k != "emb"} for o in out],
                           dtype=object),
             emb=np.stack([o["emb"] for o in out]))
    print(f"[cache] {clip}: {len(out)} embeddings written to {cache_path}",
          file=sys.stderr)
    return out


def sorted_dirs(clip):
    """Person folders you created, wherever you put them under labelling/<clip>."""
    found = {}
    for path in glob.glob(f"labelling/{clip}/**/", recursive=True):
        name = os.path.basename(path.rstrip("/"))
        if name in ("unsorted", "discard") or name.startswith("5157"):
            continue
        files = [os.path.basename(f) for f in glob.glob(path + "*.jpg")]
        if files:
            found[name] = files
    return found


# ── reference side: your labels ──────────────────────────────────────────────
ref_rows = embed_manifest(REF_CLIP)
ref_dirs = sorted_dirs(REF_CLIP)
if not ref_dirs:
    sys.exit(f"no person folders under labelling/{REF_CLIP} — sort that clip first")
file_to_person = {f: p for p, fs in ref_dirs.items() for f in fs}

ref = [(file_to_person[r["file"]], r["emb"]) for r in ref_rows
       if r["file"] in file_to_person]
people = sorted({p for p, _ in ref})
print(f"[ref] {REF_CLIP}: {len(ref)} labelled crops over {len(people)} people "
      f"{ {p: sum(1 for q, _ in ref if q == p) for p in people} }", file=sys.stderr)

R = np.stack([e for _, e in ref])
r_actor = [people.index(p) for p, _ in ref]

# The global gallery's sigmoid — NOT a fit over these four people. A Platt fit
# over a handful of identities saturates: it will hand back P=0.99 for faces it
# has no basis to separate, which is exactly how a wrong label acquires a
# convincing probability. The production fit spans the whole actor population,
# so a probability means the same thing here as it does in the matcher.
cal = sae_embed.gallery_calibration(GALLERY)
print(f"[calibration] global: {cal}  assign boundary = sim "
      f"{cal.boundary_at(ASSIGN_P):.4f}", file=sys.stderr)

# ── target side: propose ─────────────────────────────────────────────────────
tgt_rows = embed_manifest(TARGET_CLIP)
T = np.stack([t["emb"] for t in tgt_rows])
r_actor_arr = np.asarray(r_actor)
# per-actor best-of-N for every target crop at once: (n_people, n_target)
best_sim = np.stack([(R[r_actor_arr == people.index(p)] @ T.T).max(axis=0)
                     for p in people])
proposals = []
for j, t in enumerate(tgt_rows):
    k = int(np.argmax(best_sim[:, j]))
    prob = cal.probability(float(best_sim[k, j]))       # calibrated, never a bare cosine
    proposals.append({**t, "person": people[k] if prob >= ASSIGN_P else None,
                      "p": prob, "top1": people[k]})

# At the production threshold the global fit stays silent on most of these
# faces, which is the honest answer for profile and downward-gaze shots — but a
# labelling aid wants throughput, not caution. --all proposes the top-1 person
# for every crop and orders the review sheets by descending probability, so the
# proposals degrade visibly down the sheet and you can stop correcting where
# they stop being right. The probability is shown, never hidden.
if "--all" in sys.argv:
    for x in proposals:
        x["person"] = x["top1"]

named = [x for x in proposals if x["person"]]
print(f"[propose] {TARGET_CLIP}: {len(named)}/{len(proposals)} named at P>={ASSIGN_P}; "
      f"{len(proposals) - len(named)} left unsorted", file=sys.stderr)
for p in people:
    got = [x for x in named if x["person"] == p]
    if got:
        ps = [x["p"] for x in got]
        print(f"   {p}: {len(got):>3} crops   P {min(ps):.3f}–{max(ps):.3f}", file=sys.stderr)

if DRY:
    sys.exit(0)

# ── write proposed folders, mirroring the ref clip's layout ──────────────────
ref_parent = os.path.dirname(next(iter(glob.glob(f"labelling/{REF_CLIP}/**/{people[0]}/",
                                                 recursive=True))).rstrip("/"))
tgt_parent = ref_parent.replace(REF_CLIP, TARGET_CLIP)
for p in people:
    d = f"{tgt_parent}/{p}"
    if os.path.isdir(d):                # never clobber corrections already made
        print(f"[skip] {d} exists — leaving your sorting alone", file=sys.stderr)
        continue
    os.makedirs(d, exist_ok=True)
def find_crop(clip, fname):
    """Locate a crop wherever it currently sits under labelling/<clip>."""
    hits = glob.glob(f"labelling/{clip}/**/{fname}", recursive=True)
    return hits[0] if hits else None

moved = 0
for x in named:
    src = find_crop(TARGET_CLIP, x["file"])
    dst = f"{tgt_parent}/{x['person']}/{x['file']}"
    if src and os.path.abspath(src) != os.path.abspath(dst):
        shutil.move(src, dst)
        moved += 1
print(f"[write] moved {moved} crops into proposed folders", file=sys.stderr)

# ── review sheets: confirmed block, then proposed block ─────────────────────
def load(clip, person, fname):
    for cand in glob.glob(f"labelling/{clip}/**/{person}/{fname}", recursive=True):
        return cv2.imread(cand)
    return None

for person in people:
    conf = [(REF_CLIP, f, None) for f in ref_dirs.get(person, [])]
    prop = sorted([(TARGET_CLIP, x["file"], x["p"]) for x in named
                   if x["person"] == person],
                  key=lambda t: -t[2])          # most confident first
    items = conf + prop
    if not items:
        continue
    rows_n = (len(items) + SHEET_COLS - 1) // SHEET_COLS
    sheet = np.full((rows_n * (THUMB + 26), SHEET_COLS * THUMB, 3), 30, np.uint8)
    for n, (clip, fname, p) in enumerate(items):
        img = load(clip, person, fname)
        if img is None:
            continue
        rr, cc = divmod(n, SHEET_COLS)
        y, x = rr * (THUMB + 26), cc * THUMB
        sheet[y:y + THUMB, x:x + THUMB] = cv2.resize(img, (THUMB, THUMB))
        if p is None:
            tag, col = f"{clip[-3:]} CONFIRMED", (170, 170, 170)
        else:
            tag, col = f"{clip[-3:]} P={p:.2f}", (140, 255, 140)
        cv2.putText(sheet, tag, (x + 3, y + THUMB + 17),
                    cv2.FONT_HERSHEY_SIMPLEX, 0.42, col, 1)
    cv2.imwrite(f"labelling/review_{person}.jpg", sheet)
    print(f"   review_{person}.jpg: {len(conf)} confirmed + {len(prop)} proposed",
          file=sys.stderr)

json.dump({x["file"]: {"person": x["person"], "p": x["p"]} for x in proposals},
          open(f"labelling/proposed_{TARGET_CLIP}.json", "w"), indent=1)