ENH: Validate ICON finetune checkpoints and score tutorial 02 by land…

src/physiotwin4d/register_images_icon.py

@@ -11,6 +11,7 @@
 """
 
 import logging
+from pathlib import Path
 from typing import Optional, Union
 
 import itk
@@ -107,14 +108,23 @@ def set_weights_path(self, weights_path: str) -> None:
         pretrained weights. Clears any previously loaded network so the new
         weights are applied on the next call to register().
 
-        Also, use this to specify the path to store the downloaded weights.  The
-        file must not exist for the weights to be downloaded correctly.  Typical
-        suffix is ".trch".
+        The file must already exist.  uniGradICON treats a missing
+        ``weights_location`` as a download destination and silently fetches the
+        stock pretrained weights into it, so an unvalidated path yields a
+        stock-weights registration that looks like a finetuned one.
 
         Args:
-            weights_path: Path to a uniGradICON checkpoint, e.g.
-                "results/duke_4d_finetune/checkpoints/network_weights_100"
+            weights_path: Path to an existing uniGradICON checkpoint, e.g.
+                "results/duke_4d_finetune/checkpoints/network_weights_final.trch"
+
+        Raises:
+            FileNotFoundError: If weights_path does not exist.
         """
+        if not Path(weights_path).exists():
+            raise FileNotFoundError(
+                f"uniGradICON weights not found: {weights_path}.  Leave the "
+                "weights path unset to use the stock pretrained weights."
+            )
         self.weights_path = weights_path
         self.net = None  # force reload on next register() call
 

src/physiotwin4d/workflow_finetune_icon_registration.py

@@ -423,14 +423,24 @@ def expected_weights_path(self) -> Path:
         """Return the path uniGradICON writes its final checkpoint to.
 
         ``unigradicon.finetuning.finetune`` writes
-        ``<experiment.name>/checkpoints/Finetune_multi_final.trch`` at the end of
-        training.  Also the return value of :meth:`process`.
+        ``<experiment.name>/checkpoints/network_weights_final.trch`` at the end
+        of training -- its ``NETWORK_WEIGHTS_PREFIX`` plus the ``"final"`` epoch
+        label.  Also the return value of :meth:`process`.
+
+        The filename is hard-coded rather than imported: training runs in a
+        subprocess so that this process never imports ``unigradicon`` (its
+        ``finetuning`` submodule exists only on the ``feat-add-finetuning``
+        branch, so an import here would raise on a stock install).  An upstream
+        rename is caught by ``test_expected_weights_path_layout``, which
+        compares this filename against ``NETWORK_WEIGHTS_PREFIX`` wherever the
+        submodule is installed, and by the ``FileNotFoundError`` :meth:`process`
+        raises when the checkpoint is not where this says it should be.
         """
         return (
             self.experiment_dir
             / f"{self.finetune_name}_model"
             / "checkpoints"
-            / "Finetune_multi_final.trch"
+            / "network_weights_final.trch"
         )
 
     def process(self) -> Path:
@@ -441,13 +451,15 @@ def process(self) -> Path:
         existing dataset JSON or YAML in :attr:`experiment_dir` is overwritten.
 
         Returns:
-            Path to the expected final checkpoint
-            (``Finetune_multi_final.trch``).  The file is written by the
-            subprocess and exists only after a successful run.
+            Path to the final checkpoint (``network_weights_final.trch``).  The
+            file is written by the subprocess and exists only after a successful
+            run.
 
         Raises:
             subprocess.CalledProcessError: If the finetuning subprocess exits
                 with a non-zero status.
+            FileNotFoundError: If the subprocess succeeded but the checkpoint is
+                not where :meth:`expected_weights_path` says it should be.
         """
         self.log_section("FINETUNING UNIGRADICON", width=70)
 
@@ -471,6 +483,17 @@ def process(self) -> Path:
         self.log_info("Launching finetuning subprocess: %s", " ".join(cmd))
         subprocess.run(cmd, check=True, env=env)
 
+        # A missing checkpoint here is silent otherwise: uniGradICON treats an
+        # unknown weights path as a download destination, so a stale filename,
+        # or a run directory renamed with a "-N" suffix by the ``footsteps``
+        # package uniGradICON uses to lay out its runs, would yield
+        # stock-weight registrations that look finetuned.
         weights_path = self.expected_weights_path()
-        self.log_info("Finetuning complete. Expected weights at %s", weights_path)
+        if not weights_path.exists():
+            raise FileNotFoundError(
+                f"Finetuning finished but no checkpoint at {weights_path}.  "
+                "Check for a suffixed run directory (uniGradICON appends "
+                "'-1', '-2', ... when the experiment directory already exists)."
+            )
+        self.log_info("Finetuning complete. Weights at %s", weights_path)
         return weights_path

tests/test_register_images_icon.py

@@ -466,5 +466,23 @@ def test_different_iteration_counts(
         print(f"Tested {len(iteration_counts)} different iteration counts")
 
 
+def test_set_weights_path_rejects_missing_file(tmp_path: Path) -> None:
+    """A missing checkpoint must raise instead of downloading stock weights.
+
+    uniGradICON treats an unknown ``weights_location`` as a download
+    destination, so an unchecked path silently yields a stock-weights
+    registration that looks like a finetuned one.
+    """
+    registrar = RegisterImagesICON()
+    with pytest.raises(FileNotFoundError, match="weights not found"):
+        registrar.set_weights_path(str(tmp_path / "network_weights_final.trch"))
+    assert registrar.weights_path is None
+
+    existing = tmp_path / "network_weights_final.trch"
+    existing.touch()
+    registrar.set_weights_path(str(existing))
+    assert registrar.weights_path == str(existing)
+
+
 if __name__ == "__main__":
     pytest.main([__file__, "-v", "-s"])

tests/test_workflow_finetune_icon_registration.py

@@ -38,6 +38,13 @@ def _make_image(path: Path, value: int = 1) -> None:
     itk.imwrite(img, str(path), compression=True)
 
 
+def _write_fake_checkpoint(workflow: WorkflowFinetuneICONRegistration) -> None:
+    """Stand in for the checkpoint the monkey-patched subprocess never writes."""
+    weights_path = workflow.expected_weights_path()
+    weights_path.parent.mkdir(parents=True, exist_ok=True)
+    weights_path.touch()
+
+
 @pytest.fixture
 def two_subject_dataset(tmp_path: Path) -> dict[str, Any]:
     """Two patients, two frames each, with matching labelmaps and masks on disk."""
@@ -391,9 +398,17 @@ def test_expected_weights_path_layout(tmp_path: Path) -> None:
     )
     expected = workflow.expected_weights_path()
     assert expected == (
-        tmp_path / "exp" / "exp_model" / "checkpoints" / "Finetune_multi_final.trch"
+        tmp_path / "exp" / "exp_model" / "checkpoints" / "network_weights_final.trch"
     )
 
+    # The filename must track uniGradICON's own checkpoint prefix, so an
+    # upstream rename breaks this test rather than silently sending the
+    # tutorials to a path that only ever holds stock weights.  The finetuning
+    # submodule exists only on the feat-add-finetuning branch, so skip where
+    # it is absent.
+    finetune = pytest.importorskip("unigradicon.finetuning.finetune")
+    assert expected.name == f"{finetune.NETWORK_WEIGHTS_PREFIX}_final.trch"
+
 
 # ---------------------------------------------------------------------------
 # process (subprocess is monkey-patched)
@@ -407,6 +422,13 @@ def test_process_invokes_unigradicon_subprocess(
     """process launches the uniGradICON finetune module with the YAML path."""
     captured: dict[str, Any] = {}
 
+    unigradicon_src = two_subject_dataset["output_dir"].parent / "fake_unigradicon_src"
+    workflow = WorkflowFinetuneICONRegistration(
+        log_level=logging.CRITICAL,
+        unigradicon_src_path=unigradicon_src,
+        **two_subject_dataset,
+    )
+
     def fake_run(
         cmd: list[str],
         *,
@@ -416,16 +438,11 @@ def fake_run(
         captured["cmd"] = cmd
         captured["check"] = check
         captured["env"] = env
+        _write_fake_checkpoint(workflow)
         return subprocess.CompletedProcess(args=cmd, returncode=0)
 
     monkeypatch.setattr(subprocess, "run", fake_run)
 
-    unigradicon_src = two_subject_dataset["output_dir"].parent / "fake_unigradicon_src"
-    workflow = WorkflowFinetuneICONRegistration(
-        log_level=logging.CRITICAL,
-        unigradicon_src_path=unigradicon_src,
-        **two_subject_dataset,
-    )
     weights = workflow.process()
 
     assert captured["check"] is True
@@ -448,6 +465,11 @@ def test_process_without_unigradicon_src(
 ) -> None:
     """When unigradicon_src_path is None, PYTHONPATH is not prefixed."""
 
+    workflow = WorkflowFinetuneICONRegistration(
+        log_level=logging.CRITICAL,
+        **two_subject_dataset,
+    )
+
     def fake_run(
         cmd: list[str],
         *,
@@ -456,6 +478,30 @@ def fake_run(
     ) -> subprocess.CompletedProcess[bytes]:
         # No leading entry referencing a "fake" src tree.
         assert "fake_unigradicon_src" not in env.get("PYTHONPATH", "")
+        _write_fake_checkpoint(workflow)
+        return subprocess.CompletedProcess(args=cmd, returncode=0)
+
+    monkeypatch.setattr(subprocess, "run", fake_run)
+
+    workflow.process()
+
+
+def test_process_raises_when_checkpoint_missing(
+    monkeypatch: pytest.MonkeyPatch,
+    two_subject_dataset: dict[str, Any],
+) -> None:
+    """A successful subprocess that wrote no checkpoint is an error.
+
+    uniGradICON treats an unknown weights path as a download destination, so an
+    unnoticed missing checkpoint silently degrades to stock weights.
+    """
+
+    def fake_run(
+        cmd: list[str],
+        *,
+        check: bool,
+        env: dict[str, str],
+    ) -> subprocess.CompletedProcess[bytes]:
         return subprocess.CompletedProcess(args=cmd, returncode=0)
 
     monkeypatch.setattr(subprocess, "run", fake_run)
@@ -464,4 +510,5 @@ def fake_run(
         log_level=logging.CRITICAL,
         **two_subject_dataset,
     )
-    workflow.process()
+    with pytest.raises(FileNotFoundError, match="no checkpoint"):
+        workflow.process()