Codex/foreground aware mind registration

DGX_WORKFLOW.md

@@ -0,0 +1,68 @@
+# DGX validation workflow
+
+All runtime state lives under `/mnt/labstore/asharma/hack_paris`. Nothing
+uploads to Kaggle.
+
+## Environment and data
+
+```bash
+ssh dgx01
+cd /mnt/labstore/asharma/hack_paris/code
+./dgx/setup.sh
+./dgx/download_data.sh
+```
+
+`download_data.sh` accepts either an already staged competition ZIP or a
+mode-`600` credential at
+`/mnt/labstore/asharma/hack_paris/secrets/kaggle.json`. Credentials never
+belong in the repository.
+
+## Offline experiments
+
+Submit the legacy and foreground-aware pre-registered cases:
+
+```bash
+./dgx/submit_validation.sh
+```
+
+The 32-case matrix (legacy plus smoothness 2/5/10) runs at most two one-GPU
+tasks concurrently. When all tasks finish, evaluate each candidate directory:
+
+```bash
+CANDIDATE_DIR=/mnt/labstore/asharma/hack_paris/out/validation/foreground_s5 \
+  REPORT=/mnt/labstore/asharma/hack_paris/out/validation/selection-s5.json \
+  ./dgx/evaluate.sh
+```
+
+The selector writes `out/validation/selection.json` and applies the MRR gates
+defined in the implementation plan.
+
+## Competition inference
+
+Use the accepted configuration from the offline report:
+
+```bash
+sbatch --export=ALL,SPLIT=val,METHOD=accepted,EXTRA_ARGS="..." \
+  dgx/rank_pool.sbatch
+sbatch --export=ALL,SPLIT=test,METHOD=accepted,EXTRA_ARGS="..." \
+  dgx/rank_pool.sbatch
+```
+
+After both jobs finish:
+
+```bash
+METHOD=accepted ./dgx/build_candidates.sh
+```
+
+This always creates and audits a registration-only candidate. Assignment stays
+disabled unless the organizer has explicitly confirmed that every pool is a
+bijection:
+
+```bash
+BIJECTION_CONFIRMED=1 METHOD=accepted CALIBRATION=raw \
+  ./dgx/build_candidates.sh
+```
+
+No script performs a Kaggle submission. The frozen baseline remains
+`submission_deformfull.csv`, SHA-256
+`0426c138602f9fdc8c2feda1db5088ea1556a0c73253fe419afe80ce2fb7cf02`.

README.md

@@ -2,6 +2,41 @@
 
 Your task is to build a cross-modal medical image retrieval system. For each query brain MRI volume, rank all candidate target MRI volumes from the matching gallery so that the true same-subject target appears as high as possible.
 
+## Public 1.00000 Approach: Core Difference
+
+We started from Lukas Sokcevic's frozen `submission_deformfull.csv` baseline.
+That solution already used intensity normalization, affine registration,
+deformable registration, and MIND similarity. It scored public MRR `1.0` on
+Dataset 1 and Dataset 3; the remaining errors were in the independently
+deformed Dataset 2 images.
+
+We kept Lukas's Dataset 1 and Dataset 3 rankings unchanged. The improvement was
+made only to Dataset 2 by retaining the same affine-to-deformable registration
+backbone and changing its loss and final scoring:
+
+1. **Replicate-padded MIND:** replaces `torch.roll` so opposite image boundaries
+   cannot wrap together.
+2. **Foreground-aware MIND loss:** compares registered anatomy only inside the
+   valid fixed/warped foreground intersection.
+3. **Overlap penalty:** prevents a transform from receiving a good score by
+   moving anatomy outside the field of view.
+4. **Post-registration score fusion:** combines row-standardized deformable
+   MIND, NMI, and foreground overlap:
+
+```text
+score = deformable_MIND + 0.5 * NMI + 0.1 * foreground_overlap
+```
+
+The original method ranked candidates using only the final deformable MIND
+residual over the full image. The new foreground-aware score improved Dataset 2
+while preserving the already-correct Dataset 1 and Dataset 3 rankings. The
+combined submission achieved public MRR `1.00000`.
+
+Every query was ranked independently. The submitted CSV used **no Hungarian
+assignment, forced one-to-one mapping, N4 correction, manual override, or hidden
+label**. The public score covers validation queries; private test performance
+remains unknown until final evaluation.
+
 ## Kaggle Challenge
 
 https://www.kaggle.com/t/b33ec3e76c3d4e16a6b56852470b3ebf
@@ -201,6 +236,95 @@ The complete submission template contains one row for every validation and test
 - For full challenge submissions, submit both validation and test query rows in one file.
 - To focus on one dataset first, submit only that dataset's rows and omit the other datasets. The omitted datasets receive zero credit. Multiplying the displayed score by `3` gives the MRR for the submitted dataset.
 
+## Foreground-Aware MIND Registration (Public MRR 1.00000)
+
+### Starting point and attribution
+
+This work started from Lukas Sokcevic's `submission_deformfull.csv`. We froze that
+file as the baseline before running new experiments:
+
+```text
+SHA256 0426c138602f9fdc8c2feda1db5088ea1556a0c73253fe419afe80ce2fb7cf02
+```
+
+The baseline already achieved public MRR `1.0` on Dataset 1 and Dataset 3. Its
+remaining errors were in Dataset 2, where query and target volumes are deformed
+independently. We therefore kept Lukas's Dataset 1 and Dataset 3 rankings
+unchanged and replaced only the Dataset 2 rankings.
+
+### What we changed
+
+For each Dataset 2 query-target candidate pair, the new pipeline:
+
+1. normalizes each volume independently;
+2. performs affine registration followed by deformable registration;
+3. computes MIND descriptors using replicate-padded shifts instead of
+   `torch.roll`, preventing opposite image boundaries from wrapping together;
+4. creates foreground masks and evaluates MIND similarity only over valid
+   foreground overlap;
+5. penalizes low-overlap transformations so registration cannot improve its
+   score by moving anatomy outside the field of view;
+6. computes post-registration normalized mutual information (NMI) as an
+   independent structural score; and
+7. ranks every gallery independently for each query.
+
+The selected score uses row-standardized components:
+
+```text
+score = deformable_MIND + 0.5 * NMI + 0.1 * foreground_overlap
+```
+
+Affine registration remains an initialization stage, but the selected final
+score does not add the separate affine-only residual. Deformable smoothness was
+fixed at `5`.
+
+### Offline selection
+
+We evaluated the method without using hidden labels or public-leaderboard
+changes for parameter selection. Labelled Dataset 1 pairs were independently
+deformed to reproduce Dataset 2's rotations, translations, and non-linear
+displacements.
+
+| Method | Mean MRR on five development folds | MRR on 100-pair holdout | Holdout top-1 |
+| --- | ---: | ---: | ---: |
+| Frozen registration method | 0.9676 | 0.9764 | 96% |
+| Foreground-aware MIND method | **0.9926** | **0.9950** | **99%** |
+
+The holdout MRR of `0.9950` corresponds to 99 correct rank-1 matches and one
+correct rank-2 match. Because of the competition deadline, one 100-pair
+holdout was completed rather than the three originally planned.
+
+### Competition result
+
+The audited file `submission_foreground_independent.csv` contains all 377
+required rows:
+
+- Dataset 1: frozen Lukas Sokcevic baseline rankings;
+- Dataset 2: new foreground-aware MIND rankings; and
+- Dataset 3: frozen Lukas Sokcevic baseline rankings.
+
+The submission achieved public MRR `1.00000`. Kaggle computes the public score
+from validation queries and averages the three dataset MRR values. Since each
+MRR is at most `1`, an overall public score of `1.00000` means Dataset 1,
+Dataset 2, and Dataset 3 each scored `1.00000` on their public validation
+queries.
+
+This does not reveal the score on the private test queries. Their rankings are
+included in the CSV, but the private leaderboard remains unknown until final
+evaluation.
+
+### Methods explicitly not used
+
+- No Hungarian or other one-to-one assignment was applied.
+- No target was manually promoted or corrected.
+- No query-specific override was used.
+- No N4 bias-field correction was used.
+- No hidden validation or test labels were used.
+
+Every query was ranked independently from image-derived registration scores.
+The score artifact format, deterministic validation workflow, DGX jobs, and
+submission audit are documented in `DGX_WORKFLOW.md`.
+
 ## Baseline Code
 
 We provide a small MONAI + PyTorch baseline to help you get started with the challenge data format, preprocessing, training loop, and submission generation.
@@ -253,4 +377,4 @@ uv run slice_cnn_baseline.py \
   --query-csv "$DATA_ROOT/dataset1/test_queries.csv" \
   --gallery-csv "$DATA_ROOT/dataset1/test_gallery.csv" \
   --out dataset1_submission.csv
-```
+```

assemble.py

@@ -43,6 +43,18 @@ def gallery_sizes() -> dict[tuple[str, str], int]:
     }
 
 
+def gallery_ids() -> dict[tuple[str, str], set[str]]:
+    """Return the exact permitted target IDs for every retrieval pool."""
+    return {
+        (ds, split): set(
+            load_gallery(
+                DATA_ROOT, DATA_ROOT / ds / f"{split}_gallery.csv"
+            )
+        )
+        for (ds, split) in EXPECTED
+    }
+
+
 def main() -> None:
     ap = argparse.ArgumentParser()
     ap.add_argument("--out", required=True)
@@ -51,6 +63,7 @@ def main() -> None:
 
     qmap = expected_query_ids()
     gsize = gallery_sizes()
+    gids = gallery_ids()
 
     rows: dict[str, str] = {}
     for path in args.inputs:
@@ -66,9 +79,19 @@ def main() -> None:
         if qid not in rows:
             errors.append(f"missing query {qid} ({ds}/{split})")
             continue
-        n = len(rows[qid].split())
+        ranking = rows[qid].split()
+        n = len(ranking)
         if n != gsize[(ds, split)]:
             errors.append(f"{qid} ({ds}/{split}) ranks {n}, expected {gsize[(ds, split)]}")
+        if len(set(ranking)) != len(ranking):
+            errors.append(f"{qid} ({ds}/{split}) contains duplicate target IDs")
+        missing_targets = gids[(ds, split)] - set(ranking)
+        extra_targets = set(ranking) - gids[(ds, split)]
+        if missing_targets or extra_targets:
+            errors.append(
+                f"{qid} ({ds}/{split}) has wrong gallery membership "
+                f"(missing={len(missing_targets)}, extra={len(extra_targets)})"
+            )
     extra = set(rows) - set(qmap)
     if extra:
         errors.append(f"{len(extra)} unexpected query_ids (e.g. {sorted(extra)[:3]})")

assignment.py

@@ -0,0 +1,133 @@
+"""Convert score artifacts to independent or one-to-one assignment rankings."""
+from __future__ import annotations
+
+import argparse
+import csv
+from pathlib import Path
+
+import numpy as np
+from scipy.optimize import linear_sum_assignment
+
+from gpu_reg.artifacts import load_score_artifact
+
+
+def write_submission(path: Path, rows: list[dict[str, str]]) -> None:
+    path.parent.mkdir(parents=True, exist_ok=True)
+    with path.open("w", newline="") as handle:
+        writer = csv.DictWriter(
+            handle, fieldnames=["query_id", "target_id_ranking"]
+        )
+        writer.writeheader()
+        writer.writerows(rows)
+
+
+def assignment_utility(scores: np.ndarray, calibration: str) -> np.ndarray:
+    """Calibrate pairwise scores before global assignment."""
+    scores = np.asarray(scores, dtype=np.float64)
+    if calibration == "raw":
+        return scores.copy()
+    if calibration == "robust-z":
+        median = np.median(scores, axis=1, keepdims=True)
+        mad = np.median(np.abs(scores - median), axis=1, keepdims=True)
+        scale = 1.4826 * mad
+        fallback = np.std(scores, axis=1, keepdims=True)
+        scale = np.where(scale > 1e-8, scale, np.maximum(fallback, 1e-8))
+        return (scores - median) / scale
+    if calibration == "rank":
+        order = np.argsort(-scores, axis=1, kind="stable")
+        ranks = np.empty_like(order)
+        ranks[np.arange(scores.shape[0])[:, None], order] = np.arange(
+            scores.shape[1]
+        )[None, :]
+        return 1.0 / np.log2(ranks.astype(np.float64) + 2.0)
+    raise ValueError(f"unknown calibration: {calibration}")
+
+
+def assigned_columns(scores: np.ndarray, calibration: str) -> np.ndarray:
+    """Return one assigned gallery column per query."""
+    if scores.shape[0] != scores.shape[1]:
+        raise ValueError(
+            "Hungarian assignment requires equal query/gallery counts, "
+            f"got {scores.shape}"
+        )
+    utility = assignment_utility(scores, calibration)
+    row_ind, col_ind = linear_sum_assignment(-utility)
+    assigned = np.full(scores.shape[0], -1, dtype=np.int64)
+    assigned[row_ind] = col_ind
+    if np.any(assigned < 0):
+        raise RuntimeError("assignment did not cover every query")
+    return assigned
+
+
+def ranking_rows(
+    scores: np.ndarray,
+    query_ids: list[str],
+    target_ids: list[str],
+    mode: str,
+    calibration: str = "raw",
+) -> list[dict[str, str]]:
+    """Create complete rankings while changing only the assigned first choice."""
+    scores = np.asarray(scores)
+    if scores.shape != (len(query_ids), len(target_ids)):
+        raise ValueError("score matrix does not align with IDs")
+
+    base_orders = np.argsort(-scores, axis=1, kind="stable")
+    assigned = (
+        assigned_columns(scores, calibration)
+        if mode == "hungarian"
+        else base_orders[:, 0]
+    )
+    if mode not in {"independent", "hungarian"}:
+        raise ValueError(f"unknown ranking mode: {mode}")
+
+    rows: list[dict[str, str]] = []
+    for qi, qid in enumerate(query_ids):
+        first = int(assigned[qi])
+        order = [first]
+        order.extend(int(v) for v in base_orders[qi] if int(v) != first)
+        rows.append(
+            {
+                "query_id": qid,
+                "target_id_ranking": " ".join(target_ids[i] for i in order),
+            }
+        )
+    return rows
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--scores", required=True, type=Path)
+    parser.add_argument("--out", required=True, type=Path)
+    parser.add_argument(
+        "--mode", choices=["independent", "hungarian"], default="independent"
+    )
+    parser.add_argument(
+        "--calibration", choices=["raw", "robust-z", "rank"], default="raw"
+    )
+    parser.add_argument(
+        "--bijection-confirmed",
+        action="store_true",
+        help="required safety gate before emitting Hungarian competition rankings",
+    )
+    args = parser.parse_args()
+
+    if args.mode == "hungarian" and not args.bijection_confirmed:
+        raise SystemExit(
+            "Refusing Hungarian output: obtain organizer confirmation and pass "
+            "--bijection-confirmed."
+        )
+
+    artifact = load_score_artifact(args.scores)
+    rows = ranking_rows(
+        artifact["final"],
+        artifact["query_ids"],
+        artifact["target_ids"],
+        args.mode,
+        args.calibration,
+    )
+    write_submission(args.out, rows)
+    print(f"wrote {len(rows)} {args.mode} rankings -> {args.out}")
+
+
+if __name__ == "__main__":
+    main()