diff --git a/README.md b/README.md index 144b0cc..e4491eb 100644 --- a/README.md +++ b/README.md @@ -254,6 +254,100 @@ A tiny, self-contained example (two heads, including a non-contiguous id set) li --- +## Cross-photo re-identification + +Several photos of the same scene from different viewpoints, and some instance appears +in three of them — is that one instance seen three times, or three separate ones? +`infer_match()` answers that at **inference time, with no second model and nothing +retrained**: every detected instance already carries a fingerprint, the per-query +embedding that feeds the class head, and two detections are the same physical instance +when those embeddings agree. + +```python +from eomt import EoMT + +model = EoMT("runs/train/eomt-l/weights/best.pt") +result = model.infer_match("photos/subject_42/", group_by=("class",)) + +for rec in result["identities"]: + print(rec["identity_id"], rec["class_name"], rec["attributes"], rec["num_photos"]) +``` + +One call handles one folder = one subject. Each detection is stamped with an +`identity_ids` entry, `identities` summarizes each identity (dominant class, smoothed +attributes, how many photos it appears in), and `matches` records every candidate pair +— accepted or not, with the reason — so a threshold can be retuned from the JSON +without re-running inference. With `plot` you get a grid image: every photo, instances +colored by identity, plain lines linking the matched pairs (thickness = similarity, on +an absolute scale that does not vary with panel size), and a legend. + +```bash +python scripts/match.py weights/best.pt photos/subject_42/ +python scripts/match.py weights/best.pt photos/ --each-subdir # a folder per subject +``` + +### How it works + +It is the DeepSORT recipe — detect, describe, associate — applied *between photos* +instead of between video frames, which is also what [`track()`](eomt/engine/track.py) +does for video. Three EoMT properties make the fingerprint free: + +1. the model is **NMS-free**, so two overlapping instances stay two distinct queries; +2. each query owns one instance, so its embedding describes that instance alone; +3. the embedding is already computed on the way to the class head. + +Matching runs the **Hungarian** matcher per pair of photos, then links the accepted +pairs into identities with union-find. Two constraints keep that honest: one physical +instance appears **at most once per photo** (union-find would otherwise chain two +instances of the same photo together through a third), and a **merge guard** rejects a +merge whose cross-cut similarity falls below the threshold — without it, A↔B and B↔C +chain into one identity even when A and C are nothing alike. + +### Gating and thresholds + +`group_by` decides which pairs may match at all. The default `("class",)` means only +same-class instances compete. Adding attribute heads tightens it, which matters +whenever the primary class is coarser than the distinction you care about: if one class +covers instances that sit at different places on the subject, an attribute head that +separates them — `("class", "position")`, say — stops an instance in one location from +ever matching one in another, however alike they look. Any head in the checkpoint can +be named; unknown names raise before the first forward pass. `group_by=None` disables +gating entirely and needs a much higher `sim_thres`, since far more pairs then compete +with no structural prior ruling any of them out. + +`sim_thres` (default `0.7`) is the cosine floor for "same instance". Hungarian always +returns a full assignment, so the threshold is what turns *best available partner* into +*no partner — this is a new instance*. Measured on one checkpoint across four real +photo sets, with the gate on: + +| | non-candidate pairs | true-match pairs | +|---|---|---| +| mean | 0.11 – 0.16 | 0.85 – 0.98 (the true-match mode) | +| p99 / max | 0.40 – 0.67 / 0.49 – 0.79 | — | + +The default was raised from `0.6` to `0.7` after a 30-case run showed non-candidate +pairs reaching `0.73` and the two distributions overlapping in over half the cases — +`0.6` was admitting matches with no margin at all. Every run +writes its own `diagnostics` (within-identity vs across-identity percentiles) into +`_identities.json`; if those two distributions overlap, no threshold will +save the run. + +### Limitations + +- **A gate group with one instance per photo gets no benefit from the embedding.** + Hungarian has no choice to make, and only `sim_thres` can veto the pairing. The + fingerprint earns its keep where several instances of one group compete in the same + photo. +- **Small-instance recall caps what can be matched.** An instance that is never + detected in a view cannot be linked to it. +- **One folder must be one subject.** A folder holding photos of more than one subject + will happily link generic-looking instances across them; that is a data problem, not + a matching one. +- **Letterboxing shifts scale with orientation**, so a landscape and a portrait shot of + the same subject are not on quite equal footing. + +--- + ## Install ```bash @@ -296,8 +390,11 @@ For the full training recipe, every `train()` knob, and int8 compression, see th - **Pretrained COCO checkpoints.** None are published yet. COCO-trained `s`/`b`/`l` weights will be released on the Hugging Face Hub (the `from_pretrained` / `hf://` loading plumbing is already in place and waiting for them). -- **Multi-image re-ID via contrastive learning.** Train the auxiliary head with a - contrastive objective so each instance's query embedding becomes a **re-identification - vector** — matching the same object across images, frames and cameras for tracking and - retrieval. The aux head already produces a per-instance embedding from the detector's - own matched queries; re-ID reuses that signal instead of bolting on a separate model. +- **Contrastive re-ID training.** Cross-photo re-identification already works at + inference time (see [above](#cross-photo-re-identification)) on the embedding the + detector computes anyway. What remains is *training* that embedding for the job: a + contrastive objective on the matched queries would make each one a purpose-built + re-identification vector rather than a by-product of the class head, which should + widen the margin between matches and non-matches and make `sim_thres` transferable + across datasets. Feeding those embeddings into the video tracker to re-associate + objects across occlusions is the same lever. diff --git a/eomt/api.py b/eomt/api.py index 565a14a..edc7812 100644 --- a/eomt/api.py +++ b/eomt/api.py @@ -24,6 +24,7 @@ from .device import resolve_device from .engine import evaluate as _evaluate from .engine import evaluate_detection as _evaluate_detection +from .engine import match as _match from .engine import predict as _predict from .engine import track as _track from .engine import train as _train @@ -419,6 +420,46 @@ def track(self, source: str | Path, *, plot: bool = True, save: str | None = "ru hud=hud, tracker_kwargs=tracker_kwargs, **kw, ) + # ------------------------------------------------------------------ match + def infer_match(self, source: str | Path, *, plot: bool = True, + save: str | None = "runs/match", conf_thres: float = 0.3, + max_det: int = 30, mask_thresh: float = 0.5, + imgsz: int | None = None, + group_by: tuple[str, ...] | None = ("class",), + sim_thres: float = 0.7, **kw) -> dict: + """Re-identify instances across a set of photos of the *same subject*. + + Given several photos of one subject from different viewpoints, decides which + detections are the same physical instance — one instance seen from three + angles is one identity, not three. Inference only: the fingerprint is the + per-query embedding the detector already computes, so there is no second model + and nothing to retrain. + + ``source`` is one folder of photos (non-recursive) = one subject; loop for + several subjects. Each detection is stamped with an ``identity_ids`` entry, and + ``identities`` summarizes each one (dominant class, smoothed attributes, how + many photos it appears in). With ``plot`` a grid image is written under + ``save`` — every photo, instances colored by identity, connectors between + matched pairs — alongside a ``_identities.json`` summary. + + ``group_by`` gates which pairs may match at all: ``("class",)`` by default, and + any aux head name may be added — useful when the primary class is coarser than + the distinction you care about, so ``("class", "position")`` keeps instances at + different places on the subject from matching each other. ``None`` disables + gating, which needs a higher ``sim_thres``. ``sim_thres`` is the cosine floor + for "same instance" — Hungarian always returns a full assignment, so the + threshold is what turns "best available partner" into "new instance". + """ + return _match( + self.model, str(source), plot=plot, save=save, conf_thres=conf_thres, + max_det=max_det, mask_thresh=mask_thresh, + imgsz=_resolve_imgsz(imgsz, self.model), group_by=group_by, + sim_thres=sim_thres, **kw, + ) + + #: Alias so the method, the engine function and ``scripts/match.py`` agree. + match = infer_match + # ------------------------------------------------------------------- save def save(self, path: str | Path) -> None: """Write a self-describing checkpoint (reloadable with ``EoMT(path)``).""" diff --git a/eomt/engine/__init__.py b/eomt/engine/__init__.py index ff031e8..925c36f 100644 --- a/eomt/engine/__init__.py +++ b/eomt/engine/__init__.py @@ -1,8 +1,9 @@ """Training, validation and inference engines for EoMT instance segmentation.""" +from .match import match from .predict import predict from .track import track from .train import train from .validate import evaluate, evaluate_detection, sweep -__all__ = ["train", "evaluate", "evaluate_detection", "predict", "track", "sweep"] +__all__ = ["train", "evaluate", "evaluate_detection", "predict", "track", "match", "sweep"] diff --git a/eomt/engine/match.py b/eomt/engine/match.py new file mode 100644 index 0000000..64f13d4 --- /dev/null +++ b/eomt/engine/match.py @@ -0,0 +1,371 @@ +"""Cross-photo instance re-identification over a set of photos of one subject. + +Runs the still-image inference (:func:`~eomt.engine.predict.predict_image`) once per +photo, asks :mod:`eomt.reid` which detections across photos are the same physical +instance, and (optionally) renders a grid image plus a JSON summary. + +This is the photo-set counterpart to :func:`~eomt.engine.track.track`: the tracker +associates across *video frames* using motion plus appearance, this associates across +*unordered photos* using appearance alone, since there is no motion prior between two +shots taken from opposite sides of a subject. + +One call handles one set of photos of one subject. Multiple subjects are the caller's +loop — see ``scripts/match.py --each-subdir``. +""" + +from __future__ import annotations + +import json +import time +import warnings +from pathlib import Path + +import numpy as np +import torch +from PIL import Image + +from ..reid import ( + cluster, + flatten_instances, + gate_mask, + pairwise_matches, + similarity_diagnostics, + similarity_matrix, + summarize_identities, + validate_group_by, +) +from ..serialization import load_model +from ..visualize import draw_identity_grid +from .predict import _IMAGE_EXTS, _iter_sources, predict_image + + +def _panelize(image: Image.Image, result: dict, panel_size: int) -> tuple[Image.Image, dict]: + """Downscale a photo and its masks/boxes to panel scale for rendering. + + Rendering has to happen at panel scale, not full resolution: ``draw_instances`` + sizes its strokes and font from the image it is handed. Doing the downscale here + also keeps full-resolution masks from being retained for the whole run. + """ + w, h = image.size + scale = min(panel_size / max(w, h), 1.0) + nw, nh = max(1, int(round(w * scale))), max(1, int(round(h * scale))) + small = image.resize((nw, nh), Image.LANCZOS) if (nw, nh) != (w, h) else image.copy() + + out = {k: v for k, v in result.items() if k not in ("masks", "boxes", "plot", "embed")} + boxes = result["boxes"] + out["boxes"] = boxes * scale if len(boxes) else boxes + masks = result.get("masks") + if isinstance(masks, torch.Tensor) and masks.numel(): + # NEAREST, not bilinear: interpolating a bool mask and re-binarizing shifts + # the boundary and leaves a gray fringe. + arr = masks.cpu().numpy() + resized = np.stack([ + np.asarray( + Image.fromarray(m.astype(np.uint8) * 255).resize((nw, nh), Image.NEAREST) + ) > 127 + for m in arr + ]) if arr.shape[0] else np.zeros((0, nh, nw), dtype=bool) + out["masks"] = torch.from_numpy(resized) + return small, out + + +def _print_summary(identities: list[dict], num_images: int, top: int = 15) -> None: + if not identities: + return + print(f"[match] top identities by photo count (of {len(identities)}):") + for rec in identities[:top]: + attrs = " ".join(f"{k}={v}" for k, v in (rec.get("attributes") or {}).items()) + sim = rec.get("mean_similarity") + sim_txt = f" sim~{sim:.2f}" if sim is not None else "" + print( + f" #{rec['identity_id']:<3} {rec['class_name']:<18} " + f"{rec['num_photos']}/{num_images} photos{sim_txt} {attrs}" + ) + + +def match( + model, + source: str, + *, + plot: bool = True, + save: str | None = "runs/match", + conf_thres: float = 0.3, + max_det: int = 30, + mask_thresh: float = 0.5, + device: str = "auto", + imgsz: int | None = None, + group_by: tuple[str, ...] | None = ("class",), + sim_thres: float = 0.7, + center: bool = False, + guard: str = "mean", + guard_factor: float = 1.0, + keep_masks: bool = False, + panel_size: int = 480, + cols: int | None = None, + link_mode: str = "chain", + legend_attr: str | None = None, + alpha: float = 0.35, + show_scores: bool = False, +) -> dict: + """Re-identify instances across a set of photos of the same subject. + + Args: + source: a folder of photos (non-recursive) or a single image. + device: only consulted when ``model`` is a path to load; ignored when an + already-loaded model is passed (as :meth:`eomt.EoMT.infer_match` does). + group_by: gate keys — two instances may only match if they agree on all of + them. ``("class",)`` by default; add aux head names to tighten, which + matters when the primary class is coarser than the distinction you care + about (e.g. ``("class", "position")`` so instances at different places on + the subject never match, however alike they look). ``None`` disables + gating, which needs a stricter ``sim_thres`` since far more pairs then + compete. + center: mean-center embeddings before normalizing (transductive; see + :func:`eomt.reid.similarity_matrix`). + guard / guard_factor: merge guard against transitive chaining + (:func:`eomt.reid.cluster`). + keep_masks: retain full-resolution masks in the returned per-photo dicts. + Off by default — 30 photos of 20 instances at 4000x3000 is over 7 GB. + The masks are released as each photo is processed, not at the end, so + leaving this off caps peak memory rather than only trimming the result. + sim_thres: cosine floor for "same instance"; also the low end of the link + width scale, so it changes what a hairline link means. + link_mode: ``"chain"`` draws one link between panel-consecutive members of + each identity (same connectivity, far fewer lines), ``"all"`` draws every + accepted pair, ``"none"`` draws none. Links are plain lines — the "same + object" relation is symmetric, so they carry no arrowheads. + legend_attr: which attribute head the legend shows beside the class; ``None`` + uses each identity's first attribute. + + Returns: + ``{"images", "identities", "matches", "num_identities", "num_images", + "config", "diagnostics", "plot", "plot_path", "summary_path"}``. The last + three are ``None`` when the corresponding output was not produced. Each + per-photo dict gains ``identity_ids`` ``(n,)`` and ``embed`` ``(n, hidden)`` + (raw, unnormalized). Per-photo dicts hold torch tensors (as ``predict`` + does); the summaries are plain JSON-able Python (as ``track`` does). + """ + if isinstance(model, (str, Path)): + model = load_model(model, device=device) + + # Validate the gate before the first forward: a typo should cost milliseconds, + # not N ViT passes. + group_by = validate_group_by(group_by, getattr(model, "aux_specs", [])) + + dev = next(model.parameters()).device + imgsz = int(imgsz if imgsz is not None else model.image_size) + letterbox = bool(getattr(model, "preprocess_letterbox", False)) + names = getattr(model, "names", None) + + paths = list(_iter_sources(source)) + if not paths: + raise FileNotFoundError( + f"No images found in {source!r} (looked for {sorted(_IMAGE_EXTS)}). " + "Note the source is one folder of photos of one subject and is not " + "searched recursively." + ) + if len(paths) == 1: + warnings.warn( + "match() received a single photo: there is nothing to match across, so " + "every detection becomes its own identity.", + stacklevel=2, + ) + + print( + f"[match] eomt-{getattr(model, 'size', '?')} " + f"({getattr(model, 'family', 'instance')}, nc={getattr(model, 'nc', '?')}, " + f"imgsz={imgsz}) on {dev} — {len(paths)} photo(s), " + f"group_by={group_by or '(off)'}, sim_thres={sim_thres}" + ) + + # --- 1. one forward per photo ---------------------------------------- + # Each photo is reduced to panel scale and released within its own iteration: + # nothing full-resolution outlives the loop body. Holding all N photos and all N + # mask stacks to the end instead would peak at several GB on a 30-photo set of + # 4000x3000 shots — the panel is all the renderer ever needs, and ``keep_masks`` + # is what asks for the full-resolution masks to survive. + results, panels, t0 = [], [], time.perf_counter() + for i, path in enumerate(paths): + image = Image.open(path).convert("RGB") + it0 = time.perf_counter() + res = predict_image( + model, image, device=dev, imgsz=imgsz, conf_thres=conf_thres, + max_det=max_det, mask_thresh=mask_thresh, letterbox=letterbox, embed=True, + ) + if dev.type == "cuda": + torch.cuda.synchronize(dev) + res["elapsed_ms"] = (time.perf_counter() - it0) * 1e3 + res["path"] = str(path) + res["image_index"] = i + results.append(res) + if plot: + # Identity ids are not known until clustering; the panel gets them below. + small, panel_res = _panelize(image, res, panel_size) + panels.append({"image": small, "result": panel_res}) + del image + if not keep_masks: + res.pop("masks", None) + print(f"[match] {path.name}: {res['num_detections']} instances " + f"({res['elapsed_ms']:.1f} ms)") + + # --- 2. fingerprints, gate, Hungarian, link -------------------------- + dim = int(getattr(model.config, "hidden_size", 0)) + embed, photo, keys, origin = flatten_instances(results, group_by=group_by, dim=dim) + sim = similarity_matrix(embed, center=center) + gate = gate_mask(photo, keys) + pairs = pairwise_matches(sim, photo, keys, sim_thres=sim_thres) + identity, pairs = cluster( + pairs, sim, photo, num_instances=len(origin), sim_thres=sim_thres, + guard=guard, guard_factor=guard_factor, gate=gate, + ) + identities = summarize_identities( + identity, origin, results, sim, names=names, + aux_specs=getattr(model, "aux_specs", []), paths=[str(p) for p in paths], + ) + diagnostics = similarity_diagnostics(sim, gate, identity) + + # Stamp each detection with its identity, and drop heavy masks unless asked. + flat_of = {o: m for m, o in enumerate(origin)} + for i, res in enumerate(results): + n = int(res["num_detections"]) + res["identity_ids"] = torch.tensor( + [int(identity[flat_of[(i, j)]]) for j in range(n)], dtype=torch.long + ) + + accepted = [p for p in pairs if p["accepted"]] + if not accepted and len(paths) > 1: + print("[match] 0 accepted matches — try lowering sim_thres or relaxing group_by") + + matches = [ + { + "a_image": origin[p["a"]][0], "a_det": origin[p["a"]][1], + "b_image": origin[p["b"]][0], "b_det": origin[p["b"]][1], + "similarity": round(p["similarity"], 4), + "accepted": bool(p["accepted"]), + "reason": p["reason"], + "identity_id": p.get("identity_id"), + } + for p in pairs + ] + + config = { + "size": getattr(model, "size", None), + "family": getattr(model, "family", "instance"), + "nc": getattr(model, "nc", None), + "hidden": dim, + "imgsz": imgsz, + "letterbox": letterbox, + "conf_thres": conf_thres, "max_det": max_det, "mask_thresh": mask_thresh, + "group_by": list(group_by), "sim_thres": sim_thres, "center": center, + "guard": guard, "guard_factor": guard_factor, + } + + result = { + "source": str(source), + "images": results, + "identities": identities, + "matches": matches, + "num_identities": len(identities), + "num_images": len(paths), + "config": config, + "diagnostics": diagnostics, + "plot": None, + "plot_path": None, + "summary_path": None, + } + + # --- 3. render + write ------------------------------------------------ + save_dir = Path(save) if save else None + if save_dir is not None: + save_dir.mkdir(parents=True, exist_ok=True) + + if plot: + for i, (panel, res) in enumerate(zip(panels, results)): + panel["identity_ids"] = res["identity_ids"].tolist() + panel["caption"] = (f"{i}. {Path(res['path']).stem[:22]} " + f"({res['num_detections']} inst)") + mode = link_mode + if mode != "none" and len(panels) > 36: + print(f"[match] {len(panels)} panels: connectors disabled for legibility") + mode = "none" + links = _links_for(mode, accepted, origin, identity, sim) + grid = draw_identity_grid( + panels, names=names, aux_names=None, links=links, + identities=identities, cols=cols, panel_size=panel_size, + legend_attr=legend_attr, alpha=alpha, show_scores=show_scores, + # Absolute width scale: an accepted link is >= sim_thres by definition, + # so a threshold-grazing match draws hairline and a near-1.0 match draws + # thick — and the same similarity means the same thickness in every grid. + sim_range=(sim_thres, 1.0), + title=f"{Path(str(source)).name} — {len(identities)} identities " + f"across {len(paths)} photos", + ) + result["plot"] = grid + if save_dir is not None: + dst = save_dir / f"{Path(str(source)).name or 'match'}.png" + grid.save(dst) + result["plot_path"] = str(dst) + + if save_dir is not None: + summary = { + k: result[k] for k in ( + "source", "num_images", "num_identities", "config", + "diagnostics", "identities", "matches", "plot_path", + ) + } + dst = save_dir / f"{Path(str(source)).name or 'match'}_identities.json" + dst.write_text(json.dumps(summary, indent=2) + "\n") + result["summary_path"] = str(dst) + + total = time.perf_counter() - t0 + _print_summary(identities, len(paths)) + print( + f"[match] done: {len(identities)} identities from " + f"{sum(int(r['num_detections']) for r in results)} detections across " + f"{len(paths)} photos in {total:.2f} s" + + (f" -> {save_dir}" if save_dir is not None else "") + ) + return result + + +def _links_for(mode: str, accepted, origin, identity, sim) -> list[dict]: + """Connectors to draw: every accepted pair, or one chain per identity.""" + if mode == "none": + return [] + if mode == "all": + return [ + { + "a_panel": origin[p["a"]][0], "a_det": origin[p["a"]][1], + "b_panel": origin[p["b"]][0], "b_det": origin[p["b"]][1], + "similarity": p["similarity"], + "identity_id": p.get("identity_id", 0), + } + for p in accepted + ] + # "chain": one link between panel-consecutive members of each identity. Same + # connectivity, O(M) lines instead of O(M^2) — every pair of 30 photos is an + # unreadable hairball. + # A chained link is not always an accepted pair: an identity built from a-c and + # b-c puts a next to b in panel order, and that cut was never matched directly. + # Read its true cosine off the matrix rather than defaulting to 0.0, which under + # ``sim_range=(sim_thres, 1.0)`` would clamp to hairline and paint a confidently + # linked pair as a barely-made match. + s = sim.detach().cpu().numpy() + sims = {(p["a"], p["b"]): p["similarity"] for p in accepted} + sims.update({(p["b"], p["a"]): p["similarity"] for p in accepted}) + by_identity: dict[int, list[int]] = {} + for m, cid in enumerate(identity.tolist()): + by_identity.setdefault(int(cid), []).append(m) + links = [] + for cid, members in sorted(by_identity.items()): + if len(members) < 2: + continue + ordered = sorted(members, key=lambda m: origin[m]) + for a, b in zip(ordered, ordered[1:]): + links.append({ + "a_panel": origin[a][0], "a_det": origin[a][1], + "b_panel": origin[b][0], "b_det": origin[b][1], + "similarity": sims.get((a, b), float(s[a, b])), + "identity_id": cid, + }) + return links diff --git a/eomt/engine/predict.py b/eomt/engine/predict.py index c47935f..bccfd42 100644 --- a/eomt/engine/predict.py +++ b/eomt/engine/predict.py @@ -36,12 +36,20 @@ def predict_image( max_det: int = 100, mask_thresh: float = 0.5, letterbox: bool = True, + embed: bool = False, ) -> dict: """Run the model on one PIL image and return a postprocess dict. Returns a box-only :func:`~eomt.postprocess.postprocess_detection` dict for ``family="detect"`` models, else a :func:`~eomt.postprocess.postprocess_instance` dict (with masks). + + With ``embed``, the result also carries ``"embed"`` ``(N, hidden)``: the per-query + embedding of each kept detection, sliced out of the same forward pass. It is the + appearance fingerprint cross-photo re-identification matches on + (:mod:`eomt.reid`). The vectors are **raw**, not normalized — normalizing (and + optionally mean-centering) is :func:`eomt.reid.similarity_matrix`'s job, and it + needs the raw vectors to do it. """ orig_w, orig_h = image.size chw, meta = preprocess_numpy( @@ -59,14 +67,27 @@ def predict_image( if s.applies_to is not None } if getattr(model, "family", "instance") == "detect": - return postprocess_detection( + result = postprocess_detection( out, conf_thres, (orig_w, orig_h), max_det=max_det, preprocess_meta=meta, aux_scopes=aux_scopes, ) - return postprocess_instance( - out, conf_thres, (orig_w, orig_h), max_det=max_det, - mask_thresh=mask_thresh, preprocess_meta=meta, aux_scopes=aux_scopes, - ) + else: + result = postprocess_instance( + out, conf_thres, (orig_w, orig_h), max_det=max_det, + mask_thresh=mask_thresh, preprocess_meta=meta, aux_scopes=aux_scopes, + ) + if embed: + # Index with the returned ``query_idx``, never re-derive from the scores: the + # ``max_det`` topk reorders rows, so only ``query_idx`` is guaranteed aligned + # with the detections. ``.float()`` because the forward may run under autocast; + # ``.cpu()`` so a folder's worth of embeddings does not pin GPU memory. + qe = out.get("query_embed") + result["embed"] = ( + qe[0, result["query_idx"]].detach().float().cpu() + if qe is not None + else torch.zeros((result["num_detections"], 0)) + ) + return result def predict( @@ -85,6 +106,7 @@ def predict( show_scores: bool = True, color_by: str = "class", imgsz: int | None = None, + embed: bool = False, ) -> list[dict]: """Run inference on an image or a directory of images. @@ -95,6 +117,9 @@ def predict( ``path``. When ``plot`` is set, every image is rendered with masks/boxes/labels and written under ``save`` (default ``runs/predict``); the output path is added to the result dict as ``plot_path``. + + With ``embed``, each result also carries ``"embed"`` ``(N, hidden)`` — the + per-instance appearance fingerprint (see :func:`predict_image`). """ if isinstance(model, (str, Path)): model = load_model(model, device=device) @@ -123,7 +148,7 @@ def predict( result = predict_image( model, image, device=dev, imgsz=imgsz, conf_thres=conf_thres, max_det=max_det, - mask_thresh=mask_thresh, letterbox=letterbox, + mask_thresh=mask_thresh, letterbox=letterbox, embed=embed, ) if dev.type == "cuda": torch.cuda.synchronize(dev) diff --git a/eomt/model.py b/eomt/model.py index 1891840..fc9f68c 100644 --- a/eomt/model.py +++ b/eomt/model.py @@ -784,7 +784,8 @@ def __init__( self.eomt = EoMTEncoder(config, family=family) # Secondary per-instance heads (attributes). Each reads the per-query - # embedding — the input to ``class_predictor`` — captured by a hook. + # embedding — the leading ``num_queries`` rows of the encoder's final hidden + # state, i.e. exactly what ``class_predictor`` is fed. # ``aux_head_arch`` (a small MLP by default) is the shared head shape; it is # persisted in the checkpoint so reload rebuilds the same modules. self.aux_specs: list[AuxHeadSpec] = list(aux_heads or []) @@ -795,14 +796,6 @@ def __init__( for s in self.aux_specs } ) - self._query_embed: torch.Tensor | None = None - if self.aux_heads: - self.eomt.class_predictor.register_forward_hook(self._capture_query_embed) - - def _capture_query_embed(self, _module, inputs, _output): - # input to class_predictor is the per-query embedding [B, Q, hidden]; - # the last call per forward is the final layer — exactly what we want. - self._query_embed = inputs[0] def forward( self, @@ -811,7 +804,6 @@ def forward( class_labels: list[torch.Tensor] | None = None, box_labels: list[torch.Tensor] | None = None, ): - self._query_embed = None out = self.eomt( pixel_values=pixel_values, mask_labels=mask_labels, @@ -830,9 +822,13 @@ def forward( "masks_queries_logits": out.masks_queries_logits, "class_queries_logits": out.class_queries_logits, } - if self.aux_heads and self._query_embed is not None: - q = self._query_embed - result["query_embed"] = q + # The per-query embedding [B, Q, hidden]: the leading queries of the final + # hidden state, which is precisely the tensor ``class_predictor`` consumes. + # Always exposed — aux heads read it, and so does cross-photo re-identification + # on checkpoints that carry no aux heads at all. + q = out.last_hidden_state[:, : self.config.num_queries, :] + result["query_embed"] = q + if self.aux_heads: # Per-query logits over ALL queries are only consumed by postprocess at # inference; in training the aux loss applies each head to the matched # subset only, so skip the full-query application here. diff --git a/eomt/postprocess.py b/eomt/postprocess.py index 648c3b5..3fc768e 100644 --- a/eomt/postprocess.py +++ b/eomt/postprocess.py @@ -132,7 +132,8 @@ def postprocess_detection( Returns: ``{"num_detections", "boxes": (N,4) xyxy original px, "scores": (N,), - "classes": (N,)}``. If the model has secondary heads, ``output`` also carries + "classes": (N,), "query_idx": (N,)}``, where ``query_idx`` holds the kept + query indices into ``[0, Q)`` — the mapping back to per-query embeddings. If the model has secondary heads, ``output`` also carries ``"aux_queries_logits": {name: (1, Q, ns)}`` and the result gains ``"aux": {name: {"ids": (N,), "probs": (N, ns)}}`` for the same kept queries. """ @@ -151,6 +152,7 @@ def postprocess_detection( "boxes": torch.zeros((0, 4)), "scores": torch.zeros((0,)), "classes": torch.zeros((0,), dtype=torch.long), + "query_idx": torch.zeros((0,), dtype=torch.long), } if aux_logits is not None: empty["aux"] = { @@ -172,6 +174,7 @@ def postprocess_detection( "boxes": boxes, "scores": scores, "classes": classes.long(), + "query_idx": sel, } if aux_logits is not None: result["aux"] = _build_aux_result(aux_logits, sel, classes, aux_scopes) @@ -208,7 +211,8 @@ def postprocess_instance( Returns: ``{"num_detections", "boxes": (N,4), "scores": (N,), "classes": (N,), - "masks": (N,H,W)}``. If the model has secondary heads, ``output`` also + "masks": (N,H,W), "query_idx": (N,)}``, where ``query_idx`` holds the kept + query indices into ``[0, Q)`` — the mapping back to per-query embeddings. If the model has secondary heads, ``output`` also carries ``"aux_queries_logits": {name: (1, Q, ns)}`` and the result gains ``"aux": {name: {"ids": (N,), "probs": (N, ns)}}`` for the same kept queries. """ @@ -237,6 +241,7 @@ def postprocess_instance( "scores": torch.zeros((0,)), "classes": torch.zeros((0,), dtype=torch.long), "masks": torch.zeros((0, orig_h, orig_w), dtype=torch.bool), + "query_idx": torch.zeros((0,), dtype=torch.long), } if aux_logits is not None: empty["aux"] = { @@ -270,6 +275,7 @@ def postprocess_instance( "scores": scores, "classes": classes.long(), "masks": masks, + "query_idx": sel, } if aux_logits is not None: result["aux"] = _build_aux_result(aux_logits, sel, classes, aux_scopes) diff --git a/eomt/reid.py b/eomt/reid.py new file mode 100644 index 0000000..b8e5288 --- /dev/null +++ b/eomt/reid.py @@ -0,0 +1,477 @@ +"""Cross-photo instance re-identification from EoMT's own per-query embeddings. + +Given several photos of the *same subject* from different viewpoints, decide which +detected instances across those photos are the same physical instance. This is pure +inference: the appearance fingerprint is ``query_embed`` — the per-query embedding the +detector already computes on its way to the class head — so there is no second model +and nothing to retrain. + +The recipe is DeepSORT's, applied between photos instead of video frames: detect, +describe, associate with a Hungarian matcher, then link the pairwise associations into +identities. Three properties of EoMT make it work: + +* the model is **NMS-free**, so two overlapping instances stay two distinct queries; +* each query "owns" one instance, so its embedding describes that instance alone; +* embeddings are L2-normalized here, so a dot product **is** the cosine similarity. + +Everything in this module is pure: tensors in, tensors and plain dicts out. No model, +no file I/O, no PIL — :mod:`eomt.engine.match` supplies those. All instances across all +photos are flattened into one table of ``M`` rows with one ``(M, M)`` similarity matrix; +every function below indexes into that. +""" + +from __future__ import annotations + +import warnings + +import numpy as np +import torch +from scipy.optimize import linear_sum_assignment + +#: Gate key that always resolves, on any checkpoint: the primary class index. +CLASS_KEY = "class" + + +# --------------------------------------------------------------------------- +# Gating +# --------------------------------------------------------------------------- +def validate_group_by(group_by, aux_specs) -> tuple[str, ...]: + """Normalize and validate the gate keys against a checkpoint's aux heads. + + Two instances are candidates for matching only if they agree on every gate key. + ``"class"`` is always available; any aux head name may be added to tighten the + gate. That matters when the primary class is coarser than the distinction you + care about: if one class covers instances sitting at different places on the + subject, gating on ``("class", "position")`` stops an instance in one location + from ever matching one in another, however alike the two look. + + Args: + group_by: keys to gate on, or ``None`` / ``()`` to disable gating entirely. + aux_specs: the model's ``aux_specs`` (list of ``AuxHeadSpec``). + + Returns: + The normalized key tuple (possibly empty). + + Raises: + ValueError: on an unknown key — a typo must fail loudly, and it must fail + *before* the first forward pass rather than after N of them. + """ + if group_by is None: + return () + if isinstance(group_by, str): + group_by = (group_by,) + keys = tuple(str(k) for k in group_by) + if not keys: + return () + + by_name = {s.name: s for s in (aux_specs or [])} + valid = [CLASS_KEY, *by_name] + unknown = [k for k in keys if k != CLASS_KEY and k not in by_name] + if unknown: + raise ValueError( + f"Unknown group_by key(s) {unknown}. This checkpoint supports: {valid}." + ) + if len(set(keys)) != len(keys): + raise ValueError(f"Duplicate group_by key(s) in {keys}.") + + # A class-scoped head reports ``-1`` ("not applicable") for out-of-scope + # detections. Gating on it *without* "class" would make every out-of-scope + # instance compare equal, silently merging across classes. + scoped = [k for k in keys if k != CLASS_KEY and by_name[k].applies_to is not None] + if scoped and CLASS_KEY not in keys: + warnings.warn( + f"group_by={keys} gates on class-scoped head(s) {scoped} without " + f"{CLASS_KEY!r}: out-of-scope instances all share the -1 sentinel and " + "will gate together across classes. Add 'class' to group_by.", + stacklevel=2, + ) + return keys + + +def flatten_instances(results: list[dict], *, group_by: tuple[str, ...] = (CLASS_KEY,), dim: int | None = None): + """Flatten per-photo results into one instance table. + + Args: + results: per-photo dicts from :func:`eomt.engine.predict.predict_image` + called with ``embed=True``. + group_by: validated gate keys (see :func:`validate_group_by`). + dim: embedding width, used only when *every* photo is empty and there is no + row to infer it from. Pass ``model.config.hidden_size``. + + Returns: + ``(embed (M, D), photo (M,), keys, origin)`` where ``keys[m]`` is the gate + tuple for instance ``m`` (``()`` when gating is off) and ``origin[m]`` is its + ``(image_index, detection_index)``. + """ + embeds, photo, keys, origin = [], [], [], [] + for i, res in enumerate(results): + n = int(res.get("num_detections", 0)) + if n: + embeds.append(res["embed"].float()) + for j in range(n): + photo.append(i) + origin.append((i, j)) + key = [] + for k in group_by: + if k == CLASS_KEY: + key.append(int(res["classes"][j])) + else: + key.append(int(res["aux"][k]["ids"][j])) + keys.append(tuple(key)) + + if embeds: + embed = torch.cat(embeds) + else: + # No detections anywhere: the width cannot be read off the data. + embed = torch.zeros((0, int(dim or 0))) + return embed, torch.tensor(photo, dtype=torch.long), keys, origin + + +def gate_mask(photo: torch.Tensor, keys: list[tuple[int, ...]]) -> torch.Tensor: + """``(M, M)`` bool: True where a pair is a legal match candidate. + + Legal means *different photo* (one physical instance appears once per photo) and + *equal gate keys*. + + Vectorized: the key tuples are interned to ints so both conditions become + broadcast comparisons. The pairwise Python loop this replaces cost ~2.6 s at + ``M=900`` (30 photos x ``max_det=30``) — more than every other matching step + combined. + """ + m = len(keys) + if m == 0: + return torch.zeros((0, 0), dtype=torch.bool) + interned = {k: i for i, k in enumerate(dict.fromkeys(keys))} + kid = torch.tensor([interned[k] for k in keys], dtype=torch.long) + out = (photo[:, None] != photo[None, :]) & (kid[:, None] == kid[None, :]) + out.fill_diagonal_(False) # a pair is two distinct instances + return out + + +# --------------------------------------------------------------------------- +# Similarity +# --------------------------------------------------------------------------- +def l2_normalize(x: torch.Tensor, *, eps: float = 1e-8) -> torch.Tensor: + """Row-wise unit norm. A zero row stays zero rather than becoming NaN.""" + if x.numel() == 0: + return x + norm = x.norm(dim=-1, keepdim=True).clamp_min(eps) + return x / norm + + +def similarity_matrix(embed: torch.Tensor, *, center: bool = False) -> torch.Tensor: + """``(M, M)`` cosine similarity between L2-normalized embeddings. + + ``center`` subtracts the mean embedding over all instances in the set before + normalizing. That removes the common-mode component shared by every query and + slightly widens the gap between matches and non-matches, but it is *transductive* + — the similarities, and therefore the meaning of ``sim_thres``, then depend on + which photos happen to be in the set. Measured on real data the gain was small + (gap +0.44 -> +0.55 on one case) while the whole scale shifted down, so it is off + by default. + """ + if embed.numel() == 0: + return torch.zeros((embed.shape[0], embed.shape[0])) + x = embed.float() + if center: + x = x - x.mean(dim=0, keepdim=True) + z = l2_normalize(x) + return z @ z.T + + +# --------------------------------------------------------------------------- +# Pairwise matching +# --------------------------------------------------------------------------- +def pairwise_matches( + sim: torch.Tensor, + photo: torch.Tensor, + keys: list[tuple[int, ...]], + *, + sim_thres: float = 0.6, +) -> list[dict]: + """Hungarian-match every pair of photos, one gate group at a time. + + The matcher is run per ``(photo_a < photo_b)`` and, within that, per gate key + present in both — never on one masked matrix. That is not an optimization, it is + a correctness requirement: ``linear_sum_assignment`` rejects ``-inf`` outright and + raises ``cost matrix is infeasible`` on ``+inf`` whenever the smaller side cannot + be fully assigned, which is the *common* case once pairs are gated out. Because + gating is equality on a key tuple, the gate matrix is exactly block-diagonal over + gate groups, so solving each block is equivalent and no infinity ever reaches + scipy. + + Hungarian always returns a full assignment on the smaller side, so ``sim_thres`` + is what turns "best available partner" into "no partner — this is a new instance". + + Returns: + Every candidate pair, accepted or not, as + ``{"a", "b", "similarity", "accepted", "reason"}`` with global indices + ``a < b``, sorted by descending similarity. Rejections are kept so a run's + threshold can be retuned from its JSON without redoing inference. + """ + s = sim.detach().cpu().numpy() + photos = photo.tolist() + by_photo: dict[int, list[int]] = {} + for idx, p in enumerate(photos): + by_photo.setdefault(int(p), []).append(idx) + + pairs: list[dict] = [] + for pa in sorted(by_photo): + for pb in sorted(x for x in by_photo if x > pa): + # Bucket each photo's instances by gate key, then solve per shared key. + ga: dict[tuple, list[int]] = {} + for idx in by_photo[pa]: + ga.setdefault(keys[idx], []).append(idx) + gb: dict[tuple, list[int]] = {} + for idx in by_photo[pb]: + gb.setdefault(keys[idx], []).append(idx) + + for key in sorted(set(ga) & set(gb)): + rows, cols = ga[key], gb[key] + block = s[np.ix_(rows, cols)] + # float64: float32 ties resolve differently inside the solver. + ri, ci = linear_sum_assignment(-block.astype(np.float64)) + for r, c in zip(ri.tolist(), ci.tolist()): + a, b = rows[r], cols[c] + value = float(block[r, c]) + ok = value >= sim_thres + pairs.append({ + "a": a, "b": b, "similarity": value, + "accepted": ok, "reason": "" if ok else "below_thres", + }) + # Explicit tie-break: never rely on sort stability over an unordered source. + pairs.sort(key=lambda p: (-p["similarity"], p["a"], p["b"])) + return pairs + + +# --------------------------------------------------------------------------- +# Linking pairwise matches into identities +# --------------------------------------------------------------------------- +class _DisjointSet: + """Union-find carrying, per root, the member list and a photo-index bitmask.""" + + def __init__(self, n: int, photo: list[int]): + self.parent = list(range(n)) + self.rank = [0] * n + self.members = [[i] for i in range(n)] + self.photos = [1 << int(p) for p in photo] + + def find(self, x: int) -> int: + while self.parent[x] != x: + self.parent[x] = self.parent[self.parent[x]] + x = self.parent[x] + return x + + def union(self, ra: int, rb: int) -> int: + if self.rank[ra] < self.rank[rb]: + ra, rb = rb, ra + self.parent[rb] = ra + if self.rank[ra] == self.rank[rb]: + self.rank[ra] += 1 + self.members[ra] = self.members[ra] + self.members[rb] + self.photos[ra] |= self.photos[rb] + return ra + + +def cluster( + pairs: list[dict], + sim: torch.Tensor, + photo: torch.Tensor, + *, + num_instances: int, + sim_thres: float = 0.6, + guard: str = "mean", + guard_factor: float = 1.0, + gate: torch.Tensor | None = None, +) -> tuple[torch.Tensor, list[dict]]: + """Link accepted pairs into cross-photo identities with union-find. + + Accepted pairs are consumed in descending similarity, so the strongest evidence + forms clusters first and a weak pair arrives to be tested against an + already-formed cluster. The merge guard's whole meaning rests on that order, so + this function re-sorts rather than trusting the caller to have done it — with the + same explicit tie-break :func:`pairwise_matches` uses, so both agree. Two checks + run before each union: + + **Photo-uniqueness (always on).** One physical instance appears at most once per + photo, so two instances of the *same* photo must never land in one identity. Plain + union-find happily chains them via a third photo; the per-root photo bitmask makes + that structurally impossible (``reason="photo_conflict"``). + + **Merge guard** (``guard="mean" | "min" | "off"``). Union-find is transitive: A-B + and B-C chain into one identity even when A and C are nothing alike. Before + merging clusters A and B, require the similarity across the A x B cut to be at + least ``guard_factor * sim_thres`` (``reason="guard"``). Merging two singletons + skips the check — their cut is exactly the pair that already cleared the + threshold. ``"mean"`` tolerates one genuinely hard viewpoint pair inside an + otherwise strong cluster and is the default; ``"min"`` is the paranoid setting and + will block almost everything at 30 photos. + + Returns: + ``(identity (M,) long, pairs)`` — identity ids are assigned in *first + appearance* order, never by cluster size, so they stay stable (and so do the + render colors) when one detection changes between runs. ``pairs`` is a copy + with ``accepted`` / ``reason`` updated to reflect what the linker did. + """ + pairs = [dict(p) for p in pairs] + pairs.sort(key=lambda p: (-p["similarity"], p["a"], p["b"])) + ds = _DisjointSet(num_instances, photo.tolist()) + s = sim.detach().cpu().numpy() + g = gate.detach().cpu().numpy() if gate is not None else None + floor = guard_factor * sim_thres + + for p in pairs: + if not p["accepted"]: + continue + ra, rb = ds.find(p["a"]), ds.find(p["b"]) + if ra == rb: + continue # already one identity via a stronger path + if ds.photos[ra] & ds.photos[rb]: + p["accepted"], p["reason"] = False, "photo_conflict" + continue + ma, mb = ds.members[ra], ds.members[rb] + if guard != "off" and (len(ma) > 1 or len(mb) > 1): + block = s[np.ix_(ma, mb)] + if g is not None: + # Gating is transitive, so under the default gate every cut pair is + # available and this is a no-op. It stops being one the moment the + # gate becomes a non-transitive predicate (IoU, score, a learned + # compatibility) — an un-evaluable cut is not evidence for merging. + block = block[g[np.ix_(ma, mb)]] + if block.size == 0: + p["accepted"], p["reason"] = False, "guard" + continue + stat = float(block.mean()) if guard == "mean" else float(block.min()) + if stat < floor: + p["accepted"], p["reason"] = False, "guard" + continue + ds.union(ra, rb) + + # Number identities by first appearance. + identity = torch.full((num_instances,), -1, dtype=torch.long) + seen: dict[int, int] = {} + for m in range(num_instances): + root = ds.find(m) + if root not in seen: + seen[root] = len(seen) + identity[m] = seen[root] + for p in pairs: + p["identity_id"] = int(identity[p["a"]]) if p["accepted"] else None + return identity, pairs + + +# --------------------------------------------------------------------------- +# Reporting +# --------------------------------------------------------------------------- +def summarize_identities( + identity: torch.Tensor, + origin: list[tuple[int, int]], + results: list[dict], + sim: torch.Tensor, + *, + names: dict[int, str] | None = None, + aux_specs=None, + paths: list[str] | None = None, +) -> list[dict]: + """One JSON-able record per identity, sorted by how many photos it appears in. + + Attributes are smoothed across the identity's members by **mean probability**, + then argmaxed — the same running-mean idea the video tracker uses per track. + Members whose id is ``-1`` (a class-scoped head that does not apply) are excluded + from both the sum and the count rather than summed as zeros, and a head that no + member has an opinion on is omitted entirely rather than reported as class 0. + """ + s = sim.detach().cpu().numpy() + aux_names = {sp.name: sp.names for sp in (aux_specs or [])} + by_identity: dict[int, list[int]] = {} + for m, cid in enumerate(identity.tolist()): + by_identity.setdefault(int(cid), []).append(m) + + records = [] + for cid, members in by_identity.items(): + classes, photos, entries = [], [], [] + acc: dict[str, list] = {} + for m in members: + img, det = origin[m] + res = results[img] + classes.append(int(res["classes"][det])) + photos.append(img) + entries.append({ + "image_index": img, + "path": (paths[img] if paths else res.get("path", "")), + "det_index": det, + "score": float(res["scores"][det]), + "box": [round(float(v), 2) for v in res["boxes"][det].tolist()], + }) + for head, pred in (res.get("aux") or {}).items(): + hid = int(pred["ids"][det]) + if hid < 0: + continue # scoped head that does not apply: no opinion, not a vote + probs = pred["probs"][det].detach().float().cpu().numpy() + slot = acc.setdefault(head, [np.zeros_like(probs), 0]) + slot[0] += probs + slot[1] += 1 + + attribute_ids, attributes = {}, {} + for head, (total, count) in acc.items(): + if not count: + continue + hid = int(total.argmax()) + attribute_ids[head] = hid + attributes[head] = str(aux_names.get(head, {}).get(hid, hid)) + + cls = max(set(classes), key=classes.count) + cuts = [s[a, b] for i, a in enumerate(members) for b in members[i + 1:]] + records.append({ + "identity_id": int(cid), + "class": cls, + "class_name": str(names.get(cls, cls)) if names else str(cls), + "attributes": attributes, + "attribute_ids": attribute_ids, + "num_photos": len(set(photos)), + "num_instances": len(members), + "members": entries, + "mean_similarity": round(float(np.mean(cuts)), 4) if cuts else None, + "min_similarity": round(float(np.min(cuts)), 4) if cuts else None, + "max_similarity": round(float(np.max(cuts)), 4) if cuts else None, + }) + records.sort(key=lambda r: (-r["num_photos"], r["identity_id"])) + return records + + +def similarity_diagnostics( + sim: torch.Tensor, gate: torch.Tensor, identity: torch.Tensor +) -> dict: + """Percentiles of gated similarity, split within- vs across-identity. + + This is what tells you whether ``sim_thres`` is anywhere near the right value for + a dataset. If the two distributions overlap completely, the fingerprint is not + separating instances and no threshold will save the run — better to read that off + the JSON than to squint at the grid. + """ + s = sim.detach().cpu().numpy() + g = gate.detach().cpu().numpy() + ident = identity.tolist() + within, across = [], [] + for a in range(len(ident)): + for b in range(a + 1, len(ident)): + if not g[a, b]: + continue + (within if ident[a] == ident[b] else across).append(float(s[a, b])) + + def stats(v): + if not v: + return {"n": 0} + arr = np.asarray(v) + return { + "n": int(arr.size), + "mean": round(float(arr.mean()), 4), + "p10": round(float(np.percentile(arr, 10)), 4), + "p50": round(float(np.percentile(arr, 50)), 4), + "p90": round(float(np.percentile(arr, 90)), 4), + "p99": round(float(np.percentile(arr, 99)), 4), + "min": round(float(arr.min()), 4), + "max": round(float(arr.max()), 4), + } + + return {"within_identity": stats(within), "across_identity": stats(across)} diff --git a/eomt/visualize.py b/eomt/visualize.py index 6651a23..04908dd 100644 --- a/eomt/visualize.py +++ b/eomt/visualize.py @@ -3,6 +3,8 @@ from __future__ import annotations import colorsys +import math +from collections.abc import Sequence import numpy as np import torch @@ -70,6 +72,8 @@ def draw_instances( aux_multiline: bool = True, show_scores: bool = True, color_by: str = "class", + color_ids: Sequence[int] | None = None, + label_prefix: Sequence[str] | None = None, ) -> Image.Image: """Overlay masks, boxes and labels from a ``postprocess_instance`` result. @@ -89,8 +93,15 @@ def draw_instances( color_by: ``"class"`` colors masks/boxes by class index (instances of the same class share a color); ``"instance"`` colors each instance distinctly by its position, so same-class instances stand apart. + color_ids: optional per-instance color index, overriding ``color_by``. Pass + cross-photo identity ids to color each instance by *which object it is* + rather than what class it is (see :func:`draw_identity_grid`). + label_prefix: optional per-instance string prepended to the first label row, + e.g. ``"#3"`` for an identity id. """ def _color(i: int, cls: int) -> tuple[int, int, int]: + if color_ids is not None: + return class_color(int(color_ids[i])) return class_color(i if color_by == "instance" else cls) img = np.array(image.convert("RGB")).astype(np.float32) @@ -146,6 +157,8 @@ def _place(x: float, y: float, w: float, h: float) -> tuple[float, float]: text = f"{label} ({score:.2f})" if aux_multiline else f"{label} {score:.2f}" else: text = label + if label_prefix is not None: + text = f"{label_prefix[i]} {text}" aux = _aux_label(result, aux_names, i, multiline=aux_multiline, show_scores=show_scores) if aux: text += "\n" + aux @@ -159,3 +172,238 @@ def _place(x: float, y: float, w: float, h: float) -> tuple[float, float]: draw.multiline_text((lx - tb[0], ly - tb[1]), text, fill=(0, 0, 0), font=font) return out + + +# --------------------------------------------------------------------------- +# Cross-photo identity grid +# --------------------------------------------------------------------------- +def _centroid(result: dict, i: int) -> tuple[float, float] | None: + """Mask centroid for instance ``i``, falling back to the box center.""" + masks = result.get("masks") + if isinstance(masks, torch.Tensor) and i < masks.shape[0]: + m = masks[i].cpu().numpy().astype(bool) + if m.any(): + ys, xs = np.nonzero(m) + return float(xs.mean()), float(ys.mean()) + boxes = result.get("boxes") + if boxes is not None and i < len(boxes): + x1, y1, x2, y2 = (float(v) for v in boxes[i].tolist()) + return (x1 + x2) / 2, (y1 + y2) / 2 + return None + + +#: Absolute link stroke band in pixels, mapped from ``sim_range``. Deliberately a +#: constant rather than a function of ``panel_size``: the same similarity must be the +#: same thickness in every grid, whatever resolution the panels were rendered at. +LINK_WIDTH_PX = (2, 12) + + +def _connector(draw, p0, p1, color, width: int) -> None: + """A plain straight link between two instance centroids. + + No arrowheads. "These two are the same object" is a symmetric relation, so a head + at each end carried no information — and because the head was clamped to half the + connector's length, two links of equal similarity drew differently depending only + on how far apart their panels happened to sit. Stroke width alone now encodes + similarity, and it encodes nothing else. + """ + (x0, y0), (x1, y1) = p0, p1 + if math.hypot(x1 - x0, y1 - y0) < 1e-6: # both land on the same canvas point + return + draw.line([(x0, y0), (x1, y1)], fill=color, width=width) + + +def draw_identity_grid( + panels: list[dict], + *, + names: dict[int, str] | None = None, + aux_names: dict[str, dict[int, str]] | None = None, + links: list[dict] | None = None, + identities: list[dict] | None = None, + cols: int | None = None, + panel_size: int = 480, + pad: int = 12, + caption_h: int | None = None, + bg: tuple[int, int, int] = (22, 22, 26), + fg: tuple[int, int, int] = (235, 235, 235), + alpha: float = 0.35, + draw_boxes: bool = True, + show_scores: bool = False, + aux_multiline: bool = False, + link_alpha: float = 0.55, + link_width: tuple[int, int] | None = None, + sim_range: tuple[float, float] | None = None, + legend: bool = True, + legend_cols: int = 3, + legend_attr: str | None = None, + max_legend: int = 24, + title: str | None = None, +) -> Image.Image: + """Render one grid image for a set of photos linked by cross-photo identity. + + Each panel is one photo with its instances outlined and colored by identity id + (so the same physical object wears the same color in every photo it appears in), + labelled ``#id class``. Connectors are drawn between the matched instances of an + identity, so the actual pairwise links are visible and not just the final + clustering. A legend strip summarizes each identity. + + Args: + panels: one dict per photo, ``{"image": PIL.Image, "result": dict, + "identity_ids": sequence[int], "caption": str}``. Images and results must + already be at panel scale — see the note below. + draw_boxes: :func:`draw_instances` renders labels only alongside boxes, so + turning this off also hides the ``#id class`` labels the grid exists to + show. Leave it on unless you want outlines alone. + aux_names: pass to label each instance with its attributes too. ``None`` + (the default from :func:`eomt.engine.match.match`) keeps panel labels to + ``#id class``; the legend already carries the dominant attribute. + legend_attr: which attribute head to show beside the class in the legend. + ``None`` (default) uses each identity's first attribute, which is the + first head the checkpoint declares. + links: plain lines to draw between matched instances, each ``{"a_panel", + "a_det", "b_panel", "b_det", "similarity", "identity_id"}``. Stroke + thickness scales with similarity: a confident match is drawn thick, a + marginal one hairline. No arrowheads — the relation is symmetric. + link_width: ``(thinnest, thickest)`` stroke width in pixels. ``None`` uses + the absolute :data:`LINK_WIDTH_PX` band, so a given similarity draws at + the same thickness no matter what ``panel_size`` the grid was rendered + at. Pass an explicit pair only to override that deliberately. + sim_range: the ``(low, high)`` similarity band mapped onto ``link_width``. + Pass ``(sim_thres, 1.0)`` — as :func:`eomt.engine.match.match` does — to + make widths **absolute**, so the same similarity is the same thickness in + every grid and a link drawn at threshold is visibly hairline. ``None`` + falls back to min-maxing over the links actually drawn, which maximizes + contrast within one grid but means nothing across grids (and makes a lone + link thinnest regardless of how good the match is). + identities: records from :func:`eomt.reid.summarize_identities`, used for the + legend only. + + Note: + Panels must be pre-scaled by the caller. :func:`draw_instances` derives its + line width and font size from the image it is handed, so drawing on a 4000 px + photo and downscaling afterwards turns every label to mush. + + Note: + Identity colors come from :func:`class_color`, a golden-angle hue ramp. Past + roughly 20 identities adjacent hues stop being distinguishable in a small + legend swatch, so the ``#id`` text is the real key, not the color. + """ + n = len(panels) + cols = cols or max(1, math.ceil(math.sqrt(max(n, 1)))) + rows = max(1, math.ceil(n / cols)) if n else 1 + caption_h = caption_h or max(16, panel_size // 22) + cell_h = panel_size + caption_h + cap_font = _font(max(10, caption_h - 5)) + + legend_rows = 0 + swatch = 14 + legend_row_h = swatch + 10 + shown = (identities or [])[:max_legend] if legend else [] + if shown: + legend_rows = math.ceil(len(shown) / legend_cols) + ( + 1 if identities and len(identities) > max_legend else 0 + ) + title_h = (max(18, panel_size // 18) + pad) if title else 0 + legend_h = (legend_rows * legend_row_h + 2 * pad) if legend_rows else 0 + + W = cols * (panel_size + pad) + pad + H = title_h + rows * (cell_h + pad) + pad + legend_h + canvas = Image.new("RGB", (W, H), bg) + draw = ImageDraw.Draw(canvas) + + if title: + draw.text((pad, pad // 2), title, fill=fg, font=_font(max(14, panel_size // 22))) + + # --- panels ----------------------------------------------------------- + origins: list[tuple[int, int]] = [] + for k, panel in enumerate(panels): + r, c = divmod(k, cols) + ox = pad + c * (panel_size + pad) + oy = title_h + pad + r * (cell_h + pad) + image, result = panel["image"], panel["result"] + if aux_names is None: + # _aux_label falls back to raw numeric ids when it has no name map, which + # would tack an unreadable "0 0 1 4 3" row onto every label. No name map + # here means "no attribute rows", so drop the key outright. + result = {k: v for k, v in result.items() if k != "aux"} + ids = [int(v) for v in panel.get("identity_ids", [])] + rendered = draw_instances( + image, result, names=names, aux_names=aux_names, + alpha=alpha, draw_boxes=draw_boxes, aux_multiline=aux_multiline, + show_scores=show_scores, + color_ids=ids or None, + label_prefix=[f"#{i}" for i in ids] or None, + ) + # Center the (already panel-scaled) photo in its cell. + px = ox + (panel_size - rendered.width) // 2 + py = oy + (panel_size - rendered.height) // 2 + canvas.paste(rendered, (px, py)) + origins.append((px, py)) + caption = str(panel.get("caption", "")) + if caption: + draw.text((ox + 2, oy + panel_size + 2), caption, fill=fg, font=cap_font) + + # --- connectors ------------------------------------------------------- + # Drawn on an overlay and composited, so they never fully obliterate the + # instances they point at. + if links: + # Cosine similarities bunch up well below 1.0 in practice (a strong match is + # ~0.85, not ~0.99), so the top of the band is rarely reached — a generous + # span is what makes "thick = confident" legible at a glance. The band is a + # constant, NOT a function of ``panel_size``: width means similarity and + # nothing else, so rendering the same set at a different panel resolution + # must not change how thick a given match draws. + w_min, w_max = link_width or LINK_WIDTH_PX + if sim_range is not None: + lo, hi = (float(v) for v in sim_range) + else: + sims = [float(a.get("similarity", 0.0)) for a in links] + lo, hi = min(sims), max(sims) + span = (hi - lo) or 1.0 + overlay = Image.new("RGBA", canvas.size, (0, 0, 0, 0)) + odraw = ImageDraw.Draw(overlay) + for a in links: + ia, ib = int(a["a_panel"]), int(a["b_panel"]) + if not (0 <= ia < n and 0 <= ib < n): + continue + ca = _centroid(panels[ia]["result"], int(a["a_det"])) + cb = _centroid(panels[ib]["result"], int(a["b_det"])) + if ca is None or cb is None: + continue + p0 = (origins[ia][0] + ca[0], origins[ia][1] + ca[1]) + p1 = (origins[ib][0] + cb[0], origins[ib][1] + cb[1]) + norm = min(1.0, max(0.0, (float(a.get("similarity", 0.0)) - lo) / span)) + color = (*class_color(int(a.get("identity_id", 0))), int(255 * link_alpha)) + _connector(odraw, p0, p1, color, max(1, round(w_min + (w_max - w_min) * norm))) + canvas = Image.alpha_composite(canvas.convert("RGBA"), overlay).convert("RGB") + draw = ImageDraw.Draw(canvas) + + # --- legend ----------------------------------------------------------- + if shown: + lfont = _font(max(11, swatch)) + base_y = title_h + rows * (cell_h + pad) + pad + col_w = (W - 2 * pad) // max(1, legend_cols) + for i, rec in enumerate(shown): + r, c = divmod(i, legend_cols) + x = pad + c * col_w + y = base_y + r * legend_row_h + cid = int(rec["identity_id"]) + draw.rectangle([x, y, x + swatch, y + swatch], fill=class_color(cid)) + attrs = rec.get("attributes") or {} + # Which attribute to put beside the class. Defaulting to the first one + # keeps this working on any checkpoint; naming a head that does not exist + # (or that this identity has no value for) just omits it. + key = legend_attr if legend_attr is not None else next(iter(attrs), None) + value = attrs.get(key) if key is not None else None + label = f"#{cid} {rec.get('class_name', '?')}" + if value: + label += f"/{value}" + label += f" · {rec.get('num_photos', 0)}/{n} photos" + draw.text((x + swatch + 6, y), label, fill=fg, font=lfont) + if identities and len(identities) > max_legend: + y = base_y + math.ceil(len(shown) / legend_cols) * legend_row_h + draw.text( + (pad, y), f"+{len(identities) - max_legend} more identities", + fill=fg, font=lfont, + ) + return canvas diff --git a/scripts/match.py b/scripts/match.py new file mode 100644 index 0000000..7cba758 --- /dev/null +++ b/scripts/match.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python +"""Re-identify instances across a set of photos of the same subject. + +Each source folder is ONE subject photographed from several viewpoints; the script +decides which detections across those photos are the same physical instance and writes +an annotated grid plus a JSON summary. + + python scripts/match.py weights/best.pt photos/subject_42/ + python scripts/match.py weights/best.pt photos/ --each-subdir + python scripts/match.py weights/best.pt photos/subject_42/ --group-by class,position + python scripts/match.py weights/best.pt photos/subject_42/ --no-group-by --sim-thres 0.9 +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +from eomt import EoMT + + +def main() -> None: + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("weights", help="Checkpoint .pt or a run/weights folder.") + p.add_argument("source", help="Folder of photos of one subject (or a parent, with --each-subdir).") + p.add_argument("--out", default="runs/match", help="Output directory.") + p.add_argument("--conf", type=float, default=0.3, help="Detection confidence threshold.") + p.add_argument("--max-det", type=int, default=30, help="Max detections per photo.") + p.add_argument("--sim-thres", type=float, default=0.7, + help="Cosine-similarity floor for 'same instance'.") + p.add_argument("--group-by", default="class", + help="Comma-separated gate keys, e.g. class,position.") + p.add_argument("--no-group-by", action="store_true", + help="Disable gating entirely (use a higher --sim-thres).") + p.add_argument("--guard", default="mean", choices=("mean", "min", "off"), + help="Merge guard against transitive chaining.") + p.add_argument("--guard-factor", type=float, default=1.0) + p.add_argument("--center", action="store_true", + help="Mean-center embeddings before normalizing (transductive).") + p.add_argument("--links", default="chain", choices=("chain", "all", "none"), + help="Connector lines between matched instances.") + p.add_argument("--legend-attr", default=None, + help="Attribute head shown beside the class in the legend " + "(default: each identity's first attribute).") + p.add_argument("--panel-size", type=int, default=480) + p.add_argument("--each-subdir", action="store_true", + help="Treat every immediate subfolder of SOURCE as its own subject.") + p.add_argument("--no-plot", action="store_true") + p.add_argument("--device", default="auto") + args = p.parse_args() + + group_by = None if args.no_group_by else tuple( + k.strip() for k in args.group_by.split(",") if k.strip() + ) + root = Path(args.source) + if not root.exists(): + raise SystemExit(f"[match] no such source: {root}") + if args.each_subdir and not root.is_dir(): + raise SystemExit(f"[match] --each-subdir needs a parent folder, got a file: {root}") + sources = sorted(d for d in root.iterdir() if d.is_dir()) if args.each_subdir else [root] + if not sources: + raise SystemExit(f"[match] no subdirectories under {root}") + + model = EoMT(args.weights, device=args.device) + total = 0 + for src in sources: + result = model.infer_match( + src, plot=not args.no_plot, save=args.out, conf_thres=args.conf, + max_det=args.max_det, group_by=group_by, sim_thres=args.sim_thres, + guard=args.guard, guard_factor=args.guard_factor, center=args.center, + link_mode=args.links, panel_size=args.panel_size, + legend_attr=args.legend_attr, + ) + total += result["num_identities"] + print(f"[done] {total} identities across {len(sources)} subject(s) -> {args.out}") + + +if __name__ == "__main__": + main() diff --git a/tests/test_api.py b/tests/test_api.py index cb09995..761ac82 100644 --- a/tests/test_api.py +++ b/tests/test_api.py @@ -2,6 +2,8 @@ from __future__ import annotations +from pathlib import Path + import numpy as np import pytest import torch @@ -69,3 +71,109 @@ def test_checkpoint_carries_norm_metadata(tmp_path): # Reload restores the normalization onto the model. reloaded = EoMT(ckpt, device="cpu") assert tuple(round(x, 3) for x in reloaded.model.pixel_mean) == (0.485, 0.456, 0.406) + + +def _write_photos(tmp_path, sizes): + """A folder of random photos of differing sizes (exercises letterbox + panels).""" + import numpy as np + from PIL import Image + + d = tmp_path / "photos" + d.mkdir() + rng = np.random.default_rng(0) + for i, (w, h) in enumerate(sizes): + Image.fromarray(rng.integers(0, 255, (h, w, 3), dtype="uint8"), "RGB").save( + d / f"{i:02d}.png" + ) + return d + + +def test_infer_match_roundtrip(tmp_path): + import json + + from PIL import Image + + src = _write_photos(tmp_path, [(60, 80), (80, 60), (50, 50)]) + out = tmp_path / "run" + model = EoMT("s", device="cpu", nc=3, imgsz=140, pretrained=False) + result = model.infer_match( + src, plot=True, save=str(out), conf_thres=0.0, max_det=4, panel_size=96 + ) + + assert result["num_images"] == 3 and len(result["images"]) == 3 + for i, res in enumerate(result["images"]): + n = res["num_detections"] + assert res["image_index"] == i + assert len(res["query_idx"]) == len(res["embed"]) == len(res["identity_ids"]) == n + assert "masks" not in res # keep_masks defaults off + if n: + # ``embed`` is the RAW per-query embedding: normalization (and optional + # centering) belong to the similarity step, which needs the raw vectors. + assert res["embed"].dtype is torch.float32 + assert (res["embed"].norm(dim=-1) > 0).all() + from eomt.reid import l2_normalize + assert torch.allclose( + l2_normalize(res["embed"]).norm(dim=-1), torch.ones(n), atol=1e-4 + ) + + # Every identity holds at most one instance per photo. + for rec in result["identities"]: + photos = [m["image_index"] for m in rec["members"]] + assert len(photos) == len(set(photos)) + + # The summary half must be JSON-able: no tensor may have leaked into it. + summary = json.loads(Path(result["summary_path"]).read_text()) + assert summary["num_identities"] == result["num_identities"] + assert summary["config"]["group_by"] == ["class"] + assert set(summary) >= {"diagnostics", "identities", "matches", "config"} + Image.open(result["plot_path"]).close() + + +def test_infer_match_rejects_unknown_group_by(tmp_path): + src = _write_photos(tmp_path, [(40, 40), (40, 40)]) + model = EoMT("s", device="cpu", nc=3, imgsz=140, pretrained=False) + # Must raise on the typo before any forward runs (this model has no aux heads, + # which doubles as the "group_by on an aux-less checkpoint" case). + with pytest.raises(ValueError, match="Unknown group_by key"): + model.infer_match(src, plot=False, save=None, group_by=("class", "frontalty")) + # "class" alone is always valid. + model.infer_match(src, plot=False, save=None, conf_thres=0.0, max_det=2) + + +def test_chain_links_report_the_true_similarity_of_a_transitive_link(): + # An identity built from a-c and b-c puts a next to b in panel order, but that + # cut was never matched directly. Defaulting it to 0.0 would clamp to hairline + # under sim_range=(sim_thres, 1.0) and paint a confidently linked pair as a + # barely-made match; the real cosine is what gets drawn. + from eomt.engine.match import _links_for + from eomt.reid import cluster, pairwise_matches + + sim = torch.tensor([[1.0, 0.50, 0.95], [0.50, 1.0, 0.90], [0.95, 0.90, 1.0]]) + photo = torch.tensor([0, 1, 2]) + pairs = pairwise_matches(sim, photo, [(0,)] * 3, sim_thres=0.6) + identity, pairs = cluster(pairs, sim, photo, num_instances=3, sim_thres=0.6) + assert len(set(identity.tolist())) == 1 # all three are one identity + + accepted = [p for p in pairs if p["accepted"]] + assert {(p["a"], p["b"]) for p in accepted} == {(0, 2), (1, 2)} # 0-1 never was + origin = [(0, 0), (1, 0), (2, 0)] + drawn = { + (a["a_panel"], a["b_panel"]): a["similarity"] + for a in _links_for("chain", accepted, origin, identity, sim) + } + assert drawn[(0, 1)] == pytest.approx(0.50, abs=1e-5) # not 0.0 + assert drawn[(1, 2)] == pytest.approx(0.90, abs=1e-5) + + +def test_infer_match_single_photo_and_empty_source(tmp_path): + src = _write_photos(tmp_path, [(48, 48)]) + model = EoMT("s", device="cpu", nc=3, imgsz=140, pretrained=False) + with pytest.warns(UserWarning, match="single photo"): + result = model.infer_match(src, plot=False, save=None, conf_thres=0.0, max_det=3) + assert result["matches"] == [] + assert result["num_identities"] == result["images"][0]["num_detections"] + + empty = tmp_path / "empty" + empty.mkdir() + with pytest.raises(FileNotFoundError, match="No images found"): + model.infer_match(empty, plot=False, save=None) diff --git a/tests/test_reid.py b/tests/test_reid.py new file mode 100644 index 0000000..19fd991 --- /dev/null +++ b/tests/test_reid.py @@ -0,0 +1,345 @@ +"""CPU, no-model unit tests for cross-photo re-identification (:mod:`eomt.reid`). + +Everything here is driven by hand-built embeddings and synthetic result dicts, so the +matching/linking logic is tested without a forward pass anywhere. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +from eomt.reid import ( + cluster, + flatten_instances, + gate_mask, + l2_normalize, + pairwise_matches, + similarity_diagnostics, + similarity_matrix, + summarize_identities, + validate_group_by, +) + + +def _unit(*vals) -> torch.Tensor: + """A 2-D embedding on the unit circle, from angles in radians.""" + return torch.tensor([[np.cos(a), np.sin(a)] for a in vals], dtype=torch.float32) + + +def _from_sim(sim) -> torch.Tensor: + """Embeddings whose pairwise cosines reproduce ``sim`` (via its Cholesky root).""" + s = np.asarray(sim, dtype=np.float64) + w, v = np.linalg.eigh(s) + x = v @ np.diag(np.sqrt(np.clip(w, 1e-9, None))) + return torch.tensor(x, dtype=torch.float32) + + +# ---------------------------------------------------------------- similarity +def test_l2_normalize_and_similarity(): + x = torch.tensor([[3.0, 4.0], [0.0, 0.0], [1.0, 0.0]]) + z = l2_normalize(x) + assert torch.allclose(z[0].norm(), torch.tensor(1.0), atol=1e-6) + assert torch.equal(z[1], torch.zeros(2)) # zero row stays zero, never NaN + assert not torch.isnan(z).any() + + sim = similarity_matrix(x) + assert sim.shape == (3, 3) + assert torch.allclose(sim, sim.T, atol=1e-6) + assert sim[0, 0] == pytest.approx(1.0, abs=1e-5) + assert sim.min() >= -1.0001 and sim.max() <= 1.0001 + + +def test_similarity_matrix_handles_empty(): + assert similarity_matrix(torch.zeros((0, 8))).shape == (0, 0) + + +# --------------------------------------------------------------------- gating +def test_validate_group_by(): + class Spec: + def __init__(self, name, applies_to=None): + self.name = name + self.applies_to = applies_to + + specs = [Spec("color"), Spec("posture", applies_to=frozenset({0}))] + assert validate_group_by(None, specs) == () + assert validate_group_by((), specs) == () + assert validate_group_by("class", specs) == ("class",) + assert validate_group_by(("class", "color"), specs) == ("class", "color") + # "class" alone is valid even with no aux heads at all. + assert validate_group_by(("class",), []) == ("class",) + + with pytest.raises(ValueError, match="Unknown group_by key"): + validate_group_by(("class", "frontalty"), specs) + with pytest.raises(ValueError, match="Duplicate"): + validate_group_by(("class", "class"), specs) + # A class-scoped head gated on without "class": every out-of-scope instance + # shares the -1 sentinel and would gate together across classes. + with pytest.warns(UserWarning, match="sentinel"): + validate_group_by(("posture",), specs) + + +def test_gate_mask_blocks_same_photo_and_key_mismatch(): + photo = torch.tensor([0, 0, 1, 2]) + keys = [(1,), (2,), (1,), (1,)] + g = gate_mask(photo, keys) + assert not g[0, 1] # same photo + assert not g[0, 0] + assert g[0, 2] and g[2, 0] # different photo, same key + assert not g[1, 2] # key mismatch + assert g[0, 3] and g[2, 3] + assert torch.equal(g, g.T) + + +def test_flatten_instances_builds_keys_and_origin(): + results = [ + {"num_detections": 2, "embed": torch.eye(2), + "classes": torch.tensor([0, 1]), + "aux": {"side": {"ids": torch.tensor([1, 2])}}}, + {"num_detections": 1, "embed": torch.ones(1, 2), + "classes": torch.tensor([0]), + "aux": {"side": {"ids": torch.tensor([1])}}}, + ] + embed, photo, keys, origin = flatten_instances(results, group_by=("class", "side")) + assert embed.shape == (3, 2) + assert photo.tolist() == [0, 0, 1] + assert keys == [(0, 1), (1, 2), (0, 1)] + assert origin == [(0, 0), (0, 1), (1, 0)] + # Gating off -> every key is the empty tuple, so everything is a candidate. + _, _, keys_off, _ = flatten_instances(results, group_by=()) + assert keys_off == [(), (), ()] + + +def test_flatten_instances_all_empty_uses_declared_dim(): + embed, photo, keys, origin = flatten_instances( + [{"num_detections": 0}, {"num_detections": 0}], group_by=("class",), dim=384 + ) + assert embed.shape == (0, 384) # width cannot be read off the data + assert photo.numel() == 0 and keys == [] and origin == [] + + +# ------------------------------------------------------------------ matching +def test_hungarian_picks_optimal_not_greedy(): + # Greedy-by-best-cell takes (a0,b0)=.95 then is stuck with (a1,b1)=.10 -> 1.05. + # The optimal assignment is (a0,b1)+(a1,b0) = .80+.94 = 1.74. + sim = torch.tensor([ + [1.00, 0.00, 0.95, 0.80], + [0.00, 1.00, 0.94, 0.10], + [0.95, 0.94, 1.00, 0.00], + [0.80, 0.10, 0.00, 1.00], + ]) + photo = torch.tensor([0, 0, 1, 1]) + keys = [(0,)] * 4 + pairs = pairwise_matches(sim, photo, keys, sim_thres=0.0) + got = {(p["a"], p["b"]) for p in pairs} + assert got == {(0, 3), (1, 2)} + assert sum(p["similarity"] for p in pairs) == pytest.approx(1.74, abs=1e-5) + + +def test_threshold_turns_best_partner_into_new_instance(): + sim = torch.tensor([[1.0, 0.55], [0.55, 1.0]]) + photo = torch.tensor([0, 1]) + keys = [(0,), (0,)] + (pair,) = pairwise_matches(sim, photo, keys, sim_thres=0.6) + assert pair["accepted"] is False and pair["reason"] == "below_thres" + identity, _ = cluster([pair], sim, photo, num_instances=2, sim_thres=0.6) + assert identity.tolist() == [0, 1] # two separate instances + + +def test_gate_groups_never_reach_an_infeasible_hungarian(): + # Identical embeddings, different gate keys. This is the regression test for the + # -inf masking trap: scipy rejects -inf outright and raises "cost matrix is + # infeasible" on +inf when the smaller side cannot be fully assigned. Solving per + # gate group means no infinity is ever constructed. + sim = torch.ones((2, 2)) + photo = torch.tensor([0, 1]) + keys = [(0,), (1,)] + assert pairwise_matches(sim, photo, keys, sim_thres=0.0) == [] + + +def test_matches_are_sorted_and_keep_rejections(): + sim = torch.tensor([ + [1.0, 0.0, 0.9, 0.2], + [0.0, 1.0, 0.2, 0.3], + [0.9, 0.2, 1.0, 0.0], + [0.2, 0.3, 0.0, 1.0], + ]) + photo = torch.tensor([0, 0, 1, 1]) + pairs = pairwise_matches(sim, photo, [(0,)] * 4, sim_thres=0.6) + sims = [p["similarity"] for p in pairs] + assert sims == sorted(sims, reverse=True) + assert any(not p["accepted"] for p in pairs) # rejections retained for retuning + + +# ----------------------------------------------------------------- clustering +def test_guard_blocks_transitive_chain(): + # a-b and b-c both clear .6, but a and c are nothing alike. Union-find would + # happily chain all three into one identity. + sim = torch.tensor([[1.0, 0.7, 0.1], [0.7, 1.0, 0.7], [0.1, 0.7, 1.0]]) + photo = torch.tensor([0, 1, 2]) + pairs = pairwise_matches(sim, photo, [(0,)] * 3, sim_thres=0.6) + + ident_guard, out = cluster(pairs, sim, photo, num_instances=3, sim_thres=0.6, + guard="mean", guard_factor=1.0) + assert len(set(ident_guard.tolist())) == 2 # the chain is broken + assert any(p["reason"] == "guard" for p in out) + + ident_off, _ = cluster(pairs, sim, photo, num_instances=3, sim_thres=0.6, guard="off") + assert len(set(ident_off.tolist())) == 1 # and restored when the guard is off + + +def test_guard_allows_one_hard_pair_inside_a_strong_cluster(): + # Cut mean over {a,b} x {c} is (0.95 + 0.62)/2 = 0.785 >= 0.6: a genuinely hard + # viewpoint should not veto an otherwise strong identity. + sim = torch.tensor([[1.0, 0.95, 0.95], [0.95, 1.0, 0.62], [0.95, 0.62, 1.0]]) + photo = torch.tensor([0, 1, 2]) + pairs = pairwise_matches(sim, photo, [(0,)] * 3, sim_thres=0.6) + identity, _ = cluster(pairs, sim, photo, num_instances=3, sim_thres=0.6, guard="mean") + assert len(set(identity.tolist())) == 1 + # "min" is the paranoid setting and does veto it at guard_factor=1.1. + identity_min, _ = cluster(pairs, sim, photo, num_instances=3, sim_thres=0.6, + guard="min", guard_factor=1.1) + assert len(set(identity_min.tolist())) == 2 + + +def test_no_identity_holds_two_instances_of_one_photo(): + # a0 and a1 are both in photo 0; b is in photo 1 and matches both strongly. + # Plain union-find chains a0-b-a1 into one identity, which is physically + # impossible: one object appears at most once per photo. + sim = torch.tensor([[1.0, 0.99, 0.95], [0.99, 1.0, 0.96], [0.95, 0.96, 1.0]]) + photo = torch.tensor([0, 0, 1]) + pairs = [ + {"a": 0, "b": 2, "similarity": 0.95, "accepted": True, "reason": ""}, + {"a": 1, "b": 2, "similarity": 0.96, "accepted": True, "reason": ""}, + ] + identity, out = cluster(pairs, sim, photo, num_instances=3, sim_thres=0.6) + assert identity[0] != identity[1] + assert any(p["reason"] == "photo_conflict" for p in out) + # The invariant, stated globally. + for cid in set(identity.tolist()): + photos = [int(photo[m]) for m in range(3) if identity[m] == cid] + assert len(photos) == len(set(photos)) + + +def test_identity_ids_are_first_appearance_ordered(): + # Instance 0 is a singleton; 1..3 form the big cluster. Numbering by size would + # give the big cluster id 0 and flip every render color; numbering by appearance + # keeps ids stable. + sim = torch.full((4, 4), 0.9) + sim[0, :] = sim[:, 0] = 0.0 + sim[0, 0] = 1.0 + photo = torch.tensor([0, 1, 2, 3]) + pairs = pairwise_matches(sim, photo, [(0,)] * 4, sim_thres=0.6) + identity, _ = cluster(pairs, sim, photo, num_instances=4, sim_thres=0.6, guard="off") + assert identity[0].item() == 0 + assert len(set(identity[1:].tolist())) == 1 and identity[1].item() == 1 + + +def test_clustering_is_deterministic_under_ties(): + sim = torch.full((6, 6), 0.8) + sim.fill_diagonal_(1.0) + photo = torch.tensor([0, 1, 2, 3, 4, 5]) + pairs = pairwise_matches(sim, photo, [(0,)] * 6, sim_thres=0.6) + a, _ = cluster(pairs, sim, photo, num_instances=6, sim_thres=0.6) + b, _ = cluster(pairs, sim, photo, num_instances=6, sim_thres=0.6) + shuffled = list(reversed(pairs)) + c, _ = cluster(shuffled, sim, photo, num_instances=6, sim_thres=0.6) + assert torch.equal(a, b) and torch.equal(a, c) + + +def test_cluster_sorts_defensively_so_strongest_evidence_wins(): + # a0 and a1 share photo 0, so only one of them can join b — and it must be the + # stronger pair (a1-b, 0.96), whichever order the caller hands the pairs in. + # ``cluster``'s guard is meaningless unless it sees pairs strongest-first, so it + # re-sorts rather than trusting the caller to have done it. + sim = torch.tensor([[1.0, 0.99, 0.95], [0.99, 1.0, 0.96], [0.95, 0.96, 1.0]]) + photo = torch.tensor([0, 0, 1]) + weak = {"a": 0, "b": 2, "similarity": 0.95, "accepted": True, "reason": ""} + strong = {"a": 1, "b": 2, "similarity": 0.96, "accepted": True, "reason": ""} + for order in ([weak, strong], [strong, weak]): + identity, out = cluster(order, sim, photo, num_instances=3, sim_thres=0.6) + assert identity[1] == identity[2] # the 0.96 pair linked + assert identity[0] != identity[2] # the 0.95 pair lost the conflict + rejected = [p for p in out if p["reason"] == "photo_conflict"] + assert [(p["a"], p["b"]) for p in rejected] == [(0, 2)] + + +def test_degenerate_inputs(): + empty = torch.zeros((0, 0)) + ident, pairs = cluster([], empty, torch.zeros(0, dtype=torch.long), num_instances=0) + assert ident.numel() == 0 and pairs == [] + + # One photo: nothing to match across, one identity per detection. + sim = torch.eye(2) + photo = torch.tensor([0, 0]) + assert pairwise_matches(sim, photo, [(0,), (0,)], sim_thres=0.0) == [] + ident, _ = cluster([], sim, photo, num_instances=2) + assert sorted(ident.tolist()) == [0, 1] + + +# ----------------------------------------------------------------- reporting +def _result(classes, aux=None, n=None): + n = len(classes) if n is None else n + return { + "num_detections": n, + "classes": torch.tensor(classes, dtype=torch.long), + "scores": torch.ones(n), + "boxes": torch.zeros((n, 4)), + "aux": aux or {}, + "path": "p.jpg", + } + + +def test_summarize_smooths_by_mean_prob(): + # Per-member argmax is [0, 1], but the mean probability favors class 1. + aux_a = {"typ": {"ids": torch.tensor([0]), "probs": torch.tensor([[0.51, 0.49]])}} + aux_b = {"typ": {"ids": torch.tensor([1]), "probs": torch.tensor([[0.10, 0.90]])}} + results = [_result([0], aux_a), _result([0], aux_b)] + identity = torch.tensor([0, 0]) + recs = summarize_identities( + identity, [(0, 0), (1, 0)], results, torch.full((2, 2), 0.9), + names={0: "widget"}, aux_specs=[], paths=["a.jpg", "b.jpg"], + ) + assert len(recs) == 1 + assert recs[0]["attribute_ids"]["typ"] == 1 + assert recs[0]["num_photos"] == 2 and recs[0]["class_name"] == "widget" + + +def test_summarize_excludes_minus_one_members(): + # A scoped head that does not apply reports -1 with zeroed probs. Summing those + # zeros into the mean (as a naive running average would) drags the result toward + # class 0; they must be excluded from both the sum and the count. + aux_a = {"typ": {"ids": torch.tensor([1]), "probs": torch.tensor([[0.10, 0.90]])}} + aux_na = {"typ": {"ids": torch.tensor([-1]), "probs": torch.tensor([[0.0, 0.0]])}} + base = summarize_identities( + torch.tensor([0]), [(0, 0)], [_result([0], aux_a)], torch.ones((1, 1)), + names=None, aux_specs=[], paths=["a.jpg"], + ) + with_na = summarize_identities( + torch.tensor([0, 0]), [(0, 0), (1, 0)], + [_result([0], aux_a), _result([0], aux_na)], torch.full((2, 2), 0.9), + names=None, aux_specs=[], paths=["a.jpg", "b.jpg"], + ) + assert base[0]["attribute_ids"] == with_na[0]["attribute_ids"] == {"typ": 1} + + # A head no member has an opinion on is omitted, not reported as class 0. + only_na = summarize_identities( + torch.tensor([0]), [(0, 0)], [_result([0], aux_na)], torch.ones((1, 1)), + names=None, aux_specs=[], paths=["a.jpg"], + ) + assert "typ" not in only_na[0]["attributes"] + + +def test_similarity_diagnostics_splits_within_and_across(): + sim = torch.tensor([[1.0, 0.9, 0.2], [0.9, 1.0, 0.3], [0.2, 0.3, 1.0]]) + gate = torch.ones((3, 3), dtype=torch.bool) + diag = similarity_diagnostics(sim, gate, torch.tensor([0, 0, 1])) + assert diag["within_identity"]["n"] == 1 + assert diag["within_identity"]["mean"] == pytest.approx(0.9, abs=1e-4) + assert diag["across_identity"]["n"] == 2 + assert diag["across_identity"]["max"] == pytest.approx(0.3, abs=1e-4) + # An empty split reports n=0 rather than raising on an empty percentile. + assert similarity_diagnostics( + sim, torch.zeros((3, 3), dtype=torch.bool), torch.tensor([0, 1, 2]) + )["within_identity"] == {"n": 0} diff --git a/tests/test_smoke.py b/tests/test_smoke.py index 31183e6..5097567 100644 --- a/tests/test_smoke.py +++ b/tests/test_smoke.py @@ -7,6 +7,7 @@ import torch from eomt import build_eomt_config, build_model, postprocess_instance +from eomt.postprocess import postprocess_detection as _postprocess_detection from eomt.config import EOMT_CONFIGS IMGSZ = 140 # 14 * 10 -> tiny patch grid keeps the test fast @@ -34,10 +35,13 @@ def test_build_forward_shapes(): x = torch.randn(2, 3, IMGSZ, IMGSZ) with torch.no_grad(): out = model(x) - assert set(out) == {"masks_queries_logits", "class_queries_logits"} + # ``query_embed`` is emitted unconditionally — including for models with no aux + # heads, which is what lets cross-photo re-id run on any checkpoint. + assert set(out) == {"masks_queries_logits", "class_queries_logits", "query_embed"} q = EOMT_CONFIGS["s"].num_queries assert out["class_queries_logits"].shape == (2, q, NC + 1) assert out["masks_queries_logits"].shape[:2] == (2, q) + assert out["query_embed"].shape == (2, q, EOMT_CONFIGS["s"].hidden_size) def test_train_step_backward(): @@ -576,3 +580,317 @@ def test_checkpoint_stores_letterbox_mode(tmp_path): save_checkpoint(ckpt, path) loaded = load_model(path, device="cpu") assert loaded.preprocess_letterbox is True + + +def test_query_embed_matches_class_predictor_input(): + """``query_embed`` is the tensor the class head consumes — for both families. + + It is produced by slicing the encoder's final hidden state rather than by a + forward hook; this pins that the slice is the same tensor the hook used to + capture, which is what makes the hook removable. + """ + from eomt.config import EOMT_CONFIGS + + torch.manual_seed(0) + q = EOMT_CONFIGS["s"].num_queries + for family in ("instance", "detect"): + model = build_model("s", nc=NC, imgsz=IMGSZ, family=family).eval() + captured = {} + handle = model.eomt.class_predictor.register_forward_hook( + lambda _m, inputs, _o: captured.__setitem__("q", inputs[0]) + ) + with torch.no_grad(): + out = model(torch.randn(1, 3, IMGSZ, IMGSZ)) + handle.remove() + assert out["query_embed"].shape == (1, q, EOMT_CONFIGS["s"].hidden_size) + assert torch.equal(out["query_embed"], captured["q"]), family + + +def test_postprocess_returns_aligned_query_idx(): + """``query_idx`` is the detection -> query mapping the whole re-id feature rests on.""" + torch.manual_seed(0) + q, cls = 40, 3 + logits = torch.randn(1, q, cls + 1) + out = { + "masks_queries_logits": torch.randn(1, q, 8, 8), + "class_queries_logits": logits, + } + res = postprocess_instance(out, 0.0, (16, 12), max_det=10, mask_thresh=0.5) + sel = res["query_idx"] + assert sel.dtype is torch.int64 + assert sel.shape == (10,) == res["classes"].shape # the topk path fired + assert len(set(sel.tolist())) == 10 + assert sel.min() >= 0 and sel.max() < q + # The alignment guarantee: recomputing the class from the raw logits at the + # kept query indices must reproduce exactly what postprocess returned. + recomputed = logits[0].softmax(-1)[:, :-1].max(-1).indices[sel] + assert torch.equal(recomputed, res["classes"]) + + +def test_query_idx_empty_branches(): + torch.manual_seed(0) + q = 8 + out = { + "masks_queries_logits": torch.randn(1, q, 8, 8), + "class_queries_logits": torch.randn(1, q, NC + 1), + "pred_boxes": torch.rand(1, q, 4), + } + for fn, kwargs in ( + (postprocess_instance, {"mask_thresh": 0.5}), + (_postprocess_detection, {}), + ): + res = fn(out, 1.1, (16, 12), max_det=5, **kwargs) # nothing clears conf=1.1 + assert res["num_detections"] == 0 + assert res["query_idx"].shape == (0,) and res["query_idx"].dtype is torch.int64 + + +def test_query_idx_survives_min_mask_area(): + torch.manual_seed(0) + q = 20 + out = { + "masks_queries_logits": torch.randn(1, q, 8, 8) - 0.5, + "class_queries_logits": torch.randn(1, q, NC + 1), + } + res = postprocess_instance( + out, 0.0, (16, 12), max_det=q, mask_thresh=0.5, min_mask_area=20.0 + ) + n = res["num_detections"] + assert len(res["query_idx"]) == n == len(res["classes"]) == len(res["masks"]) + + +def test_predict_image_embed_is_aligned_and_unit_norm(): + from PIL import Image + + from eomt.config import EOMT_CONFIGS + from eomt.engine.predict import predict_image + + torch.manual_seed(0) + model = build_model("s", nc=NC, imgsz=IMGSZ).eval() + image = Image.fromarray( + (torch.rand(40, 60, 3).numpy() * 255).astype("uint8"), mode="RGB" + ) + res = predict_image( + model, image, device=torch.device("cpu"), imgsz=IMGSZ, + conf_thres=0.0, max_det=5, embed=True, + ) + assert res["embed"].shape == (res["num_detections"], EOMT_CONFIGS["s"].hidden_size) + assert len(res["embed"]) == len(res["query_idx"]) + assert res["embed"].dtype is torch.float32 + + +def test_draw_instances_color_ids_and_label_prefix_are_additive(): + from PIL import Image + + from eomt.visualize import draw_instances + + torch.manual_seed(0) + image = Image.new("RGB", (40, 30), (10, 10, 10)) + result = { + "num_detections": 1, + "boxes": torch.tensor([[4.0, 4.0, 20.0, 18.0]]), + "scores": torch.tensor([0.9]), + "classes": torch.tensor([1]), + "masks": (torch.rand(1, 30, 40) > 0.5), + } + base = draw_instances(image, result, {1: "widget"}) + # Defaults left alone -> byte-identical to the pre-change path. + assert np.array_equal(np.asarray(base), np.asarray(draw_instances(image, result, {1: "widget"}))) + # A different color index paints differently. + a = draw_instances(image, result, {1: "widget"}, color_ids=[5]) + b = draw_instances(image, result, {1: "widget"}, color_ids=[6]) + assert not np.array_equal(np.asarray(a), np.asarray(b)) + # A label prefix changes the rendering too. + p = draw_instances(image, result, {1: "widget"}, label_prefix=["#7"]) + assert not np.array_equal(np.asarray(p), np.asarray(base)) + + +def test_draw_identity_grid_geometry_and_degenerates(): + from PIL import Image + + from eomt.visualize import draw_identity_grid + + torch.manual_seed(0) + + def panel(n): + return { + "image": Image.new("RGB", (60, 40), (30, 30, 30)), + "result": { + "num_detections": n, + "boxes": torch.tensor([[2.0, 2.0, 30.0, 25.0]] * n).reshape(n, 4), + "scores": torch.ones(n), + "classes": torch.zeros(n, dtype=torch.long), + "masks": (torch.rand(n, 40, 60) > 0.5), + }, + "identity_ids": list(range(n)), + "caption": "0. photo", + } + + identities = [{"identity_id": 0, "class_name": "widget", "attributes": {}, + "num_photos": 2}] + links = [{"a_panel": 0, "a_det": 0, "b_panel": 1, "b_det": 0, + "similarity": 0.8, "identity_id": 0}] + grid = draw_identity_grid( + [panel(1), panel(1)], names={0: "widget"}, links=links, + identities=identities, cols=2, panel_size=64, title="t", + ) + assert grid.mode == "RGB" + assert grid.width == 2 * (64 + 12) + 12 + assert np.asarray(grid).std() > 0 # something was actually drawn + + # A photo with no detections still gets its panel, and no links is fine. + grid2 = draw_identity_grid([panel(0), panel(1)], cols=2, panel_size=64, links=[]) + assert grid2.size[0] == 2 * (64 + 12) + 12 + + +def test_draw_identity_grid_legend_attr_is_head_agnostic(): + """The legend must not hardcode any one checkpoint's attribute head name.""" + from PIL import Image + + from eomt.visualize import draw_identity_grid + + panels = [{ + "image": Image.new("RGB", (40, 40), (30, 30, 30)), + "result": { + "num_detections": 0, + "boxes": torch.zeros((0, 4)), + "scores": torch.zeros(0), + "classes": torch.zeros(0, dtype=torch.long), + }, + "identity_ids": [], + "caption": "0. photo", + }] + identities = [{ + "identity_id": 0, "class_name": "widget", "num_photos": 1, + "attributes": {"alpha": "first", "beta": "second"}, + }] + + # panel_size drives the canvas width; too narrow and every legend string clips + # at the same character, making distinct labels compare equal. + def render(**kw): + return np.asarray( + draw_identity_grid(panels, identities=identities, cols=1, panel_size=320, **kw) + ) + + default = render() # -> first head ("alpha") + explicit_first = render(legend_attr="alpha") + explicit_second = render(legend_attr="beta") + assert np.array_equal(default, explicit_first) + assert not np.array_equal(explicit_first, explicit_second) + + # A head no identity carries just omits the attribute instead of raising. + missing = render(legend_attr="not_a_head") + assert not np.array_equal(missing, explicit_first) + # ...and matches an identity that has no attributes at all. + bare = np.asarray(draw_identity_grid( + panels, cols=1, panel_size=320, + identities=[{"identity_id": 0, "class_name": "widget", "num_photos": 1, + "attributes": {}}], + )) + assert np.array_equal(missing, bare) + + +def test_draw_identity_grid_link_width_tracks_similarity(): + """A confident link must draw thicker than a marginal one, on an absolute scale.""" + from PIL import Image + + from eomt.visualize import draw_identity_grid + + torch.manual_seed(0) + + def panel(): + m = torch.zeros(1, 40, 60, dtype=torch.bool) + m[0, 18:22, 28:32] = True # a small, centered blob -> a stable centroid + return { + "image": Image.new("RGB", (60, 40), (0, 0, 0)), + "result": { + "num_detections": 1, + "boxes": torch.tensor([[28.0, 18.0, 32.0, 22.0]]), + "scores": torch.ones(1), + "classes": torch.zeros(1, dtype=torch.long), + "masks": m, + }, + "identity_ids": [0], + "caption": "", + } + + def ink(sim, **kw): + """Pixels the connector overlay painted, i.e. how heavy the stroke is.""" + grid = draw_identity_grid( + [panel(), panel()], cols=2, panel_size=240, legend=False, + draw_boxes=False, alpha=0.0, + links=[{"a_panel": 0, "a_det": 0, "b_panel": 1, "b_det": 0, + "similarity": sim, "identity_id": 0}], + **kw, + ) + return int((np.asarray(grid).sum(axis=2) > 0).sum()) + + band = {"sim_range": (0.6, 1.0)} + assert ink(0.98, **band) > ink(0.80, **band) > ink(0.61, **band) + + # Absolute, not relative: a lone link's width depends only on its similarity, + # so a near-perfect match never renders hairline just because it is alone. + assert ink(0.99, **band) > ink(0.62, **band) + + # Out-of-band similarities clamp instead of overshooting the width range. + assert ink(2.0, **band) == ink(1.0, **band) + assert ink(-1.0, **band) == ink(0.6, **band) + + # An explicit width range is honored. + assert ink(0.9, sim_range=(0.6, 1.0), link_width=(1, 2)) < ink( + 0.9, sim_range=(0.6, 1.0), link_width=(8, 16) + ) + + +def test_link_width_is_independent_of_panel_size(): + """The same similarity must draw the same stroke at any panel resolution. + + Width encodes similarity and nothing else, so a grid rendered at 240 px and the + same grid at 720 px must put an identical number of pixels into the link. Deriving + the band from ``panel_size`` (as this once did) made a match look weaker simply + because the grid was rendered smaller. + """ + from PIL import Image + + from eomt.visualize import LINK_WIDTH_PX, draw_identity_grid + + def panel(size): + h = w = size + m = torch.zeros(1, h, w, dtype=torch.bool) + m[0, h // 2 - 2:h // 2 + 2, w // 2 - 2:w // 2 + 2] = True + return { + "image": Image.new("RGB", (w, h), (0, 0, 0)), + "result": { + "num_detections": 1, + "boxes": torch.tensor([[w / 2 - 2, h / 2 - 2, w / 2 + 2, h / 2 + 2]]), + "scores": torch.ones(1), + "classes": torch.zeros(1, dtype=torch.long), + "masks": m, + }, + "identity_ids": [0], + "caption": "", + } + + def stroke_px(panel_size, pad=12): + """Link thickness, isolated by diffing against a link-free render. + + Measured in the gutter between the two panels, where nothing else is drawn, + so panel content and the canvas background cancel out. Total ink would scale + with the link's *length* (panels sit farther apart in a bigger grid), which + is why this samples one column instead. + """ + kw = dict(cols=2, panel_size=panel_size, legend=False, draw_boxes=False, + alpha=0.0, sim_range=(0.7, 1.0), pad=pad) + bare = np.asarray(draw_identity_grid( + [panel(panel_size), panel(panel_size)], links=[], **kw)).astype(int) + with_link = np.asarray(draw_identity_grid( + [panel(panel_size), panel(panel_size)], + links=[{"a_panel": 0, "a_det": 0, "b_panel": 1, "b_det": 0, + "similarity": 0.85, "identity_id": 0}], **kw)).astype(int) + gutter_x = pad + panel_size + pad // 2 + changed = (np.abs(with_link - bare).sum(axis=2) > 0)[:, gutter_x] + return int(changed.sum()) + + widths = {p: stroke_px(p) for p in (240, 480, 720)} + assert len(set(widths.values())) == 1, f"stroke varies with panel_size: {widths}" + lo, hi = LINK_WIDTH_PX + assert lo <= next(iter(widths.values())) <= hi