Give every kept scan a version from each model the site offers

flask-server/.env.example

@@ -258,6 +258,6 @@ BODYMAPS_LOCAL_AUTO_SIGNIN=false
 # scripts/user_dataset_second_opinions.py runs every kept scan through the other models, only while the site is
 # idle and a GPU worker is free, and never on this host's GPU. All settings below are optional; see
 # deploy/USER_DATASET_SECOND_OPINIONS.md.
-# SECOND_OPINIONS_MODELS=ePAI,Atlas-Net,LesionSegmenter
+# SECOND_OPINIONS_MODELS=ePAI,Atlas-Net,LesionSegmenter,SuPreM,R-Super-Pancreas-Kidney,R-Super-Merlin
 # SECOND_OPINIONS_MAX_RUN_SECONDS=1800
 # SECOND_OPINIONS_DAILY_CAP=100

flask-server/deploy/USER_DATASET_SECOND_OPINIONS.md

@@ -2,8 +2,8 @@
 
 Every scan kept in `UserData` carries the mask of the model the user ran. This
 process runs the same stored CT through the **other** models and saves each
-result beside the original, so one scan has an ePAI version, an Atlas-Net
-version and a LesionSegmenter version.
+result beside the original, so one scan has a version from every model the site offers
+(ePAI, Atlas-Net, LesionSegmenter, SuPreM, R-Super-Pancreas-Kidney, R-Super-Merlin).
 
 It is **low priority**: it only works while nobody is using the site and a GPU
 worker is idle, it never uses bdmap1's GPU, and it stops the moment a user job
@@ -112,7 +112,7 @@ failure it gives up (listed by `--status`); delete its entry in
 
 | Variable | Default | Meaning |
 | --- | --- | --- |
-| `SECOND_OPINIONS_MODELS` | `ePAI,Atlas-Net,LesionSegmenter` | models every scan should have |
+| `SECOND_OPINIONS_MODELS` | `ePAI,Atlas-Net,LesionSegmenter,SuPreM,R-Super-Pancreas-Kidney,R-Super-Merlin` | models every scan should have |
 | `SECOND_OPINIONS_BASE_URL` | `http://127.0.0.1:8000` | the website, for `/api/health` |
 | `SECOND_OPINIONS_IDLE_CHECKS` / `_CHECK_INTERVAL` | 3 / 20 s | consecutive idle samples needed |
 | `SECOND_OPINIONS_POLL_SECONDS` | 3 | how often to look for users while a model runs |
@@ -124,6 +124,11 @@ failure it gives up (listed by `--status`); delete its entry in
 
 ## Notes
 
+* **Adding a model to the list** needs it to run on the workers, to write `combined_labels.nii.gz` in the viewer's label
+  scheme (all six do), and a checkpoint rule in `services/model_provenance.py`. The R-Super releases have no label map
+  in `MODEL_TO_VIEWER`, so their `label_support` says `known: false` rather than guess which labels they can produce.
+* A bigger model list means more GPU runs per scan; `SECOND_OPINIONS_DAILY_CAP` (100 per rolling 24 hours) is what paces
+  the backlog of scans already kept, so raise it only when the workers are known to be idle.
 * The stored scans are **not de-identified** (each case keeps `user_id` and
   `session_id`), so these versions are internal until that is decided.
 * The workers must hold the same model environments and weights as the website

flask-server/scripts/user_dataset_second_opinions.py

@@ -46,7 +46,7 @@ def main(argv=None) -> int:
     mode.add_argument("--once", action="store_true")
     mode.add_argument("--loop", action="store_true")
     parser.add_argument("--limit", type=int, default=None, help="stop after this many finished tasks")
-    parser.add_argument("--models", default=None, help="comma-separated models (default: ePAI,Atlas-Net,LesionSegmenter)")
+    parser.add_argument("--models", default=None, help="comma-separated models (default: every model the site offers)")
     parser.add_argument("--root", default=None, help="UserData folder (default: USER_DATASET_PATH)")
     parser.add_argument("--json", action="store_true", help="print --status/--dry-run as JSON")
     args = parser.parse_args(argv)

flask-server/services/model_provenance.py

@@ -37,6 +37,22 @@
     "LesionSegmenter": ("LESIONSEG_CKPT_PATH",
                         "/home/visitor/lesionsegmenter/model/nnUNet_results/Dataset_LesionSegmenter/nnUNetTrainer__nnUNetPlans__3d_fullres"),
 }
+# Models whose weights are one file, not an nnU-Net folder: model -> (checkpoint env var, default path,
+# class-list env var or None, default class list or None). Same names and defaults as the runners in
+# services/auto_segmentor.py (SuPreM, RSUPER_VARIANTS); a test keeps them equal.
+MODEL_CHECKPOINT_FILE = {
+    "SuPreM": ("SUPREM_CHECKPOINT_PATH",
+               "/home/visitor/suprem_native/workspace/SuPreM/pretrained_checkpoints/supervised_suprem_unet_2100.pth",
+               None, None),
+    "R-Super-Pancreas-Kidney": ("RSUPER_PANKID_CHECKPOINT_PATH",
+                                "/home/visitor/rsuper/R-SuperPancreasKidney/R_Super_Kidney_Pancreas_UCSF_Atlas3/fold_0_latest.pth",
+                                "RSUPER_PANKID_CLASS_LIST",
+                                "/home/visitor/rsuper/R-SuperPancreasKidney/labels_rsuper_mask.yaml"),
+    "R-Super-Merlin": ("RSUPER_MERLIN_CHECKPOINT_PATH",
+                       "/home/visitor/rsuper/Merlin-Cancer-Super/checkpoint/fold_0_latest.pth",
+                       "RSUPER_MERLIN_CLASS_LIST",
+                       "/home/visitor/rsuper/Merlin-Cancer-Super/checkpoint/label_names.yaml"),
+}
 DEFAULT_CHECKPOINT_NAME = "checkpoint_final.pth"
 # Files of an nnU-Net model folder that define the model (weights + preprocessing/architecture plan).
 _PLAN_FILES = ("plans.json", "dataset.json")
@@ -83,10 +99,36 @@ def checkpoint_name(model: str, env=None) -> str:
     return DEFAULT_CHECKPOINT_NAME
 
 
+def _file_checkpoint_fingerprint(model: str, env) -> dict:
+    """A model whose weights are one file: hash it, and its class list when it has one (that list
+    decides which label each output channel gets, so it is part of what made the labels)."""
+    ckpt_var, ckpt_default, list_var, list_default = MODEL_CHECKPOINT_FILE[model]
+    path = env.get(ckpt_var, ckpt_default)
+    out: dict = {"path": path, "hashed": False}
+    if not os.path.isfile(path):
+        out["reason"] = "checkpoint file not found on this host"
+        return out
+    try:
+        files = {os.path.basename(path): _cached_sha256(path)}
+        if list_var:
+            class_list = env.get(list_var, list_default)
+            if os.path.isfile(class_list):
+                files[os.path.basename(class_list)] = _cached_sha256(class_list)
+            else:
+                out["class_list_missing"] = class_list
+    except OSError as e:
+        out["reason"] = f"unreadable:{type(e).__name__}"
+        return out
+    out.update(hashed=True, files=files)
+    return out
+
+
 def checkpoint_fingerprint(model: str, env=None) -> dict:
     """Identify the weights this process would load for `model`. Always returns a dict; `hashed` says
     whether the files could be read here."""
     env = os.environ if env is None else env
+    if model in MODEL_CHECKPOINT_FILE:
+        return _file_checkpoint_fingerprint(model, env)
     rule = MODEL_CHECKPOINT.get(model)
     if rule is None:
         return {"hashed": False, "reason": "no checkpoint rule for this model"}
@@ -148,6 +190,13 @@ def inference_settings(model: str, env=None) -> Optional[dict]:
                 # the recall-oriented operating point applied when the labels are exported
                 "lesion_threshold": float(env.get("LESIONSEG_LESION_THRESHOLD", "0.05") or "0.05"),
                 "lesion_channel": int(env.get("LESIONSEG_LESION_CHANNEL", "40") or "40")}
+    if model == "SuPreM":
+        # Fixed on the runner's command line: the U-Net backbone with the SuPreM weights.
+        return {"backbone": "unet", "suprem": True, "checkpoint_name": os.path.basename(MODEL_CHECKPOINT_FILE[model][1])}
+    if model in ("R-Super-Pancreas-Kidney", "R-Super-Merlin"):
+        # RSUPER_VARIANTS: the Pancreas-Kidney release masks lesions with its own organ masks, Merlin does not.
+        return {"checkpoint_name": os.path.basename(MODEL_CHECKPOINT_FILE[model][1]),
+                "organ_mask_on_lesion": model == "R-Super-Pancreas-Kidney"}
     return None
 
 

flask-server/services/user_dataset_models.py

@@ -3,7 +3,7 @@
 A scan is kept in UserData with the mask of the model the user ran. This module (driven by
 scripts/user_dataset_second_opinions.py, a separate process from the website) runs the same stored CT through
 the remaining models and saves each result beside the original, so one scan carries an ePAI version, an
-Atlas-Net version and a LesionSegmenter version:
+Atlas-Net version, a LesionSegmenter version and so on for every model the site offers:
 
     mask_only/USER_00000017/
         combined_labels.nii.gz, segmentations/, metadata.json     the model the user ran (unchanged)
@@ -45,7 +45,8 @@
 from services.atomic_write import atomic_destination
 
 SCHEMA = 1
-DEFAULT_MODELS = ("ePAI", "Atlas-Net", "LesionSegmenter")
+# Every model the site offers, so each kept scan ends up with one version from each.
+DEFAULT_MODELS = ("ePAI", "Atlas-Net", "LesionSegmenter", "SuPreM", "R-Super-Pancreas-Kidney", "R-Super-Merlin")
 BY_MODEL = "by_model"
 STATE_FILE = "second_opinions_state.json"
 AUDIT_FILE = "second_opinions.jsonl"

flask-server/tests/unit/test_model_provenance.py

@@ -76,7 +76,7 @@ def test_settings_follow_the_environment():
                  "checkpoint_name": "checkpoint_best.pth", "tensorrt_plan": "/p/plan.trt"}
     s = mp.inference_settings("LesionSegmenter", env={"LESIONSEG_LESION_THRESHOLD": "0.2", "LESIONSEG_LESION_CHANNEL": "41"})
     assert s["lesion_threshold"] == 0.2 and s["lesion_channel"] == 41
-    assert mp.inference_settings("SuPreM", env={}) is None
+    assert mp.inference_settings("OpenVAE", env={}) is None             # no rule for a model that is not a segmenter here
 
 
 # ----------------------------- checkpoint identity -----------------------------
@@ -133,7 +133,7 @@ def test_checkpoint_that_cannot_be_read_says_why(tmp_path):
     assert fp["hashed"] is False and "checkpoint_final.pth not found" in fp["reason"]
     fp = mp.checkpoint_fingerprint("ePAI", env={"EPAI_CKPT_PATH": str(empty), "EPAI_CHECKPOINT_NAME": "best.pth"})
     assert "best.pth not found" in fp["reason"]                              # honours the configured file name
-    assert mp.checkpoint_fingerprint("SuPreM", env={}) == {"hashed": False, "reason": "no checkpoint rule for this model"}
+    assert mp.checkpoint_fingerprint("OpenVAE", env={}) == {"hashed": False, "reason": "no checkpoint rule for this model"}
 
 
 def test_epai_tensorrt_engine_is_fingerprinted_when_configured(tmp_path):
@@ -220,3 +220,65 @@ def boom(*a, **k):
     out = mp.describe("LesionSegmenter", real_names(), (22,))
     assert out["provenance"] == {"error": "RuntimeError: disk on fire"}
     assert out["label_support"]["known"] is True
+
+
+# ----------------------------- models whose weights are a single file -----------------------------
+FILE_MODELS = ["SuPreM", "R-Super-Pancreas-Kidney", "R-Super-Merlin"]
+
+
+def test_single_file_rules_match_the_runners():
+    """Provenance repeats the runners' environment variables and defaults; a stale copy would hash the wrong file."""
+    from services.auto_segmentor import RSUPER_VARIANTS
+    for model in ("R-Super-Pancreas-Kidney", "R-Super-Merlin"):
+        v = RSUPER_VARIANTS[model]
+        ckpt_var, ckpt_default, list_var, list_default = mp.MODEL_CHECKPOINT_FILE[model]
+        assert (ckpt_var, ckpt_default) == (v["checkpoint_env"], v["checkpoint"])
+        assert (list_var, list_default) == (v["class_list_env"], v["class_list"])
+        assert mp.inference_settings(model, env={})["organ_mask_on_lesion"] is v["organ_mask_on_lesion"]
+    var, default, _, _ = mp.MODEL_CHECKPOINT_FILE["SuPreM"]
+    body = _function_source(SEGMENTOR, "_run_suprem_inference")
+    assert re.search(rf'"{var}",\s*"{re.escape(default)}"', body), "_run_suprem_inference no longer reads this checkpoint"
+    assert '"--backbone", "unet"' in body and '"--suprem"' in body       # the fixed settings recorded for SuPreM
+
+
+@pytest.mark.parametrize("model", FILE_MODELS)
+def test_a_single_file_checkpoint_is_hashed_with_its_class_list(model, tmp_path):
+    ckpt_var, _, list_var, _ = mp.MODEL_CHECKPOINT_FILE[model]
+    weights = tmp_path / "fold_0_latest.pth"; weights.write_bytes(b"weights-" + model.encode())
+    env = {ckpt_var: str(weights)}
+    if list_var:
+        classes = tmp_path / "labels.yaml"; classes.write_bytes(b"liver: 1\n")
+        env[list_var] = str(classes)
+    fp = mp.checkpoint_fingerprint(model, env=env)
+    assert fp["hashed"] is True and fp["path"] == str(weights)
+    assert fp["files"]["fold_0_latest.pth"]["sha256"] == hashlib.sha256(b"weights-" + model.encode()).hexdigest()
+    assert fp["files"]["fold_0_latest.pth"]["size_bytes"] == len(b"weights-" + model.encode())
+    if list_var:
+        assert fp["files"]["labels.yaml"]["sha256"] == hashlib.sha256(b"liver: 1\n").hexdigest()
+
+
+@pytest.mark.parametrize("model", FILE_MODELS)
+def test_a_missing_single_file_checkpoint_says_why(model, tmp_path):
+    ckpt_var = mp.MODEL_CHECKPOINT_FILE[model][0]
+    fp = mp.checkpoint_fingerprint(model, env={ckpt_var: str(tmp_path / "absent.pth")})
+    assert fp["hashed"] is False and fp["reason"] == "checkpoint file not found on this host"
+
+
+def test_a_missing_class_list_is_noted_but_the_weights_are_still_hashed(tmp_path):
+    weights = tmp_path / "w.pth"; weights.write_bytes(b"w")
+    fp = mp.checkpoint_fingerprint("R-Super-Merlin", env={"RSUPER_MERLIN_CHECKPOINT_PATH": str(weights),
+                                                           "RSUPER_MERLIN_CLASS_LIST": str(tmp_path / "gone.yaml")})
+    assert fp["hashed"] is True and list(fp["files"]) == ["w.pth"] and fp["class_list_missing"].endswith("gone.yaml")
+
+
+def test_settings_for_the_single_file_models_are_fixed_by_the_runners():
+    assert mp.inference_settings("SuPreM", env={}) == {"backbone": "unet", "suprem": True,
+                                                       "checkpoint_name": "supervised_suprem_unet_2100.pth"}
+    assert mp.inference_settings("R-Super-Pancreas-Kidney", env={})["organ_mask_on_lesion"] is True
+    assert mp.inference_settings("R-Super-Merlin", env={})["organ_mask_on_lesion"] is False
+
+
+def test_label_support_for_the_r_super_releases_says_it_is_unknown_rather_than_guess():
+    for model in ("R-Super-Pancreas-Kidney", "R-Super-Merlin"):
+        assert mp.label_support(model, real_names())["known"] is False
+    assert mp.label_support("SuPreM", real_names())["known"] is True

flask-server/tests/unit/test_user_dataset_models.py

@@ -8,6 +8,7 @@
 import http.server
 import importlib
 import json
+import re
 import os
 import sys
 import threading
@@ -112,14 +113,17 @@ def deps(self):
                       local_hostname=lambda: "bdmap1", now=self.now, sleep=self.sleep)
 
 
+THREE = ("ePAI", "Atlas-Net", "LesionSegmenter")   # the loop tests use three models; DEFAULT_MODELS has its own tests
+
+
 @pytest.fixture()
 def world(tmp_path):
     root = tmp_path / "UserData"
     sessions = tmp_path / "tmp"
     sessions.mkdir()
     return types.SimpleNamespace(root=str(root), sessions=str(sessions), rec=Recorder(str(root)),
                                  cfg=M.Config(idle_checks=2, check_interval=0, poll_seconds=0.02, max_run_seconds=30,
-                                              loop_sleep=0, idle_sleep=0))
+                                              loop_sleep=0, idle_sleep=0, models=THREE))
 
 
 def run(world, **kw):
@@ -135,17 +139,17 @@ def test_pending_tasks_skip_the_users_own_model_finished_versions_and_backoff(wo
     make_case(world.root, "USER_00000001", "ePAI")
     make_case(world.root, "USER_00000002", "LesionSegmenter")
     state = M.load_state(world.root)
-    tasks = M.pending_tasks(world.root, M.DEFAULT_MODELS, state, 1000.0)
+    tasks = M.pending_tasks(world.root, THREE, state, 1000.0)
     assert [(t.case_id, t.model) for t in tasks] == [
         ("USER_00000001", "Atlas-Net"), ("USER_00000001", "LesionSegmenter"),
         ("USER_00000002", "ePAI"), ("USER_00000002", "Atlas-Net")]
     assert all(t.model != t.origin for t in tasks)
     M.note_result(state, tasks[0], "failed", 1000.0, "boom")
-    assert tasks[0].key not in {t.key for t in M.pending_tasks(world.root, M.DEFAULT_MODELS, state, 1001.0)}
-    assert tasks[0].key in {t.key for t in M.pending_tasks(world.root, M.DEFAULT_MODELS, state, 1000.0 + 3601)}
+    assert tasks[0].key not in {t.key for t in M.pending_tasks(world.root, THREE, state, 1001.0)}
+    assert tasks[0].key in {t.key for t in M.pending_tasks(world.root, THREE, state, 1000.0 + 3601)}
     # a model the user ran that is not in the configured set still gets every configured model
     make_case(world.root, "USER_00000003", "SuPreM")
-    assert len([t for t in M.pending_tasks(world.root, M.DEFAULT_MODELS, M.load_state(world.root), 0) if t.case_id == "USER_00000003"]) == 3
+    assert len([t for t in M.pending_tasks(world.root, THREE, M.load_state(world.root), 0) if t.case_id == "USER_00000003"]) == 3
 
 
 def test_failures_back_off_then_give_up_and_preemption_is_not_a_failure():
@@ -371,7 +375,7 @@ def test_the_whole_backlog_is_worked_through_once_and_never_repeated(world):
     assert counts["ok"] == 6 and counts["failed"] == 0 and len(world.rec.calls) == 6
     for i, model in enumerate(["ePAI", "LesionSegmenter", "Atlas-Net"], 1):
         meta = json.loads(Path(world.root, "mask_only", f"USER_0000000{i}", "metadata.json").read_text())
-        assert sorted(v["model"] for v in meta["model_versions"]) == sorted(M.DEFAULT_MODELS)
+        assert sorted(v["model"] for v in meta["model_versions"]) == sorted(THREE)
         assert meta["model_versions"][0] == {"model": model, "role": "user_run", "mask": ".", "pipeline": {"segmentation_input": "original"}}
     assert os.listdir(world.sessions) == []
     again = run(world, once=True)
@@ -519,6 +523,34 @@ def test_the_safety_relies_on_these_behaviours_of_the_runner_and_the_worker_modu
     # a worker run happens only for a session folder under .../tmp/ with an id-shaped name: ours qualifies
     session = "/home/visitor/PanTS-Viewer/tmp/so-USER_00000017-LesionSegmenter-3fa9c2d1"
     assert gpu_workers._is_session_path(session)
-    assert M._SAFE_MODEL.match("Atlas-Net") and all(gpu_workers._SESSION_ID.match(f"so-USER_00000017-{m}-3fa9c2d1") for m in M.DEFAULT_MODELS)
+    assert all(M._SAFE_MODEL.match(m) for m in M.DEFAULT_MODELS) and all(gpu_workers._SESSION_ID.match(f"so-USER_00000017-{m}-3fa9c2d1") for m in M.DEFAULT_MODELS)
     # every model run here must be one the runner knows, and the user dataset must record the same names
-    assert all(f"model == '{m}'" in seg for m in M.DEFAULT_MODELS)
+    # (the R-Super releases are matched as a group: `model in RSUPER_VARIANTS`)
+    assert all(f"model == '{m}'" in seg or m in auto_segmentor.RSUPER_VARIANTS for m in M.DEFAULT_MODELS)
+
+
+# ------------------------------------------------------------------ the default model list
+def test_the_default_models_are_every_model_the_site_offers():
+    """A model added to the picker must also be added here, or scans kept afterwards would never get its version."""
+    page = (Path(__file__).resolve().parents[3] / "PanTS-Demo" / "src" / "routes" / "UploadPage.tsx").read_text(encoding="utf-8")
+    block = page[page.index("const MODEL_OPTIONS"): page.index("const MODEL_OPTIONS") + 12000]
+    offered = [m for m in re.findall(r'^    id: "([^"]+)"', block, re.M) if m != "None"]
+    assert offered and set(offered) == set(M.DEFAULT_MODELS), (offered, M.DEFAULT_MODELS)
+    assert len(set(M.DEFAULT_MODELS)) == len(M.DEFAULT_MODELS)
+    assert M.Config.from_env({}).models == M.DEFAULT_MODELS
+
+
+def test_every_default_model_has_a_checkpoint_rule_and_settings():
+    from services import model_provenance as mp
+    for model in M.DEFAULT_MODELS:
+        assert model in mp.MODEL_CHECKPOINT or model in mp.MODEL_CHECKPOINT_FILE, model
+        assert mp.inference_settings(model, env={}) is not None, model
+
+
+def test_a_kept_scan_is_owed_a_version_from_every_other_model(world):
+    make_case(world.root, "USER_00000001", "SuPreM")
+    make_case(world.root, "USER_00000002", "R-Super-Merlin")
+    tasks = M.pending_tasks(world.root, M.DEFAULT_MODELS, M.load_state(world.root), 0)
+    by_case = {c: sorted(t.model for t in tasks if t.case_id == c) for c in ("USER_00000001", "USER_00000002")}
+    assert by_case["USER_00000001"] == sorted(set(M.DEFAULT_MODELS) - {"SuPreM"})
+    assert by_case["USER_00000002"] == sorted(set(M.DEFAULT_MODELS) - {"R-Super-Merlin"})