diff --git a/docs/assets/tutorial_02_finetuning.png b/docs/assets/tutorial_02_finetuning.png index 71574a1..3ecddbd 100644 --- a/docs/assets/tutorial_02_finetuning.png +++ b/docs/assets/tutorial_02_finetuning.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:27b37d28ffd9cb7c08d4fb8649c51d3288c82054d0bd77dc2ea6329cf41b9b12 -size 4015 +oid sha256:1d0647ee4648d5abdd9a15d250e692677622d30f14a3a5b0384c6d3981b4c97c +size 17425 diff --git a/docs/assets/tutorial_05_heart_vtk_to_usd.png b/docs/assets/tutorial_05_heart_vtk_to_usd.png index 5fb6640..b80b343 100644 --- a/docs/assets/tutorial_05_heart_vtk_to_usd.png +++ b/docs/assets/tutorial_05_heart_vtk_to_usd.png @@ -1,3 +1,3 @@ version https://git-lfs.github.com/spec/v1 -oid sha256:3fc515afbb08947c6ab23261ca2bd95d7926e3d2eef3b833959a3823d1fb3836 -size 244389 +oid sha256:0ab8a4365d86d1652f1fc61ed141c8c3def30559428ef6cbda7eaaac67a0446f +size 200941 diff --git a/src/physiotwin4d/register_images_icon.py b/src/physiotwin4d/register_images_icon.py index 74e10cd..1c96408 100644 --- a/src/physiotwin4d/register_images_icon.py +++ b/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 diff --git a/src/physiotwin4d/workflow_finetune_icon_registration.py b/src/physiotwin4d/workflow_finetune_icon_registration.py index 13c78a1..42600dd 100644 --- a/src/physiotwin4d/workflow_finetune_icon_registration.py +++ b/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 - ``/checkpoints/Finetune_multi_final.trch`` at the end of - training. Also the return value of :meth:`process`. + ``/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 diff --git a/tests/test_register_images_icon.py b/tests/test_register_images_icon.py index 8b25431..7ecde32 100644 --- a/tests/test_register_images_icon.py +++ b/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"]) diff --git a/tests/test_workflow_finetune_icon_registration.py b/tests/test_workflow_finetune_icon_registration.py index 1bc3d79..e6ff5b3 100644 --- a/tests/test_workflow_finetune_icon_registration.py +++ b/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() diff --git a/tutorials/tutorial_02_lung_finetune_icon.py b/tutorials/tutorial_02_lung_finetune_icon.py index 9d8b939..0d04de6 100644 --- a/tutorials/tutorial_02_lung_finetune_icon.py +++ b/tutorials/tutorial_02_lung_finetune_icon.py @@ -4,29 +4,32 @@ Purpose ------- Finetune uniGradICON on every DIR-Lab 4D CT case except Case 1, then register -``Case1Pack_T30.mha`` (moving) to ``Case1Pack_T70.mha`` (fixed) three ways: -``RegisterImagesGreedy`` alone with its default settings, and -``RegisterImagesGreedyICON`` with the stock uniGradICON weights and with the -finetuned weights. Case 1 is never seen during finetuning, so it is a held-out +``Case1Pack_T00.mha`` (moving) to ``Case1Pack_T50.mha`` (fixed) three ways: +``RegisterImagesGreedy`` alone, deformable, with its default iteration +schedule, and ``RegisterImagesICON`` with the stock uniGradICON weights and +with the finetuned weights. Case 1 is never seen during finetuning, so it is a held-out evaluation pair. -Accuracy is measured by label overlap. ``SegmentNVSegmentCTMRI`` segments the -fixed image, and segments each registered moving image after it is warped onto -the fixed grid; every labelmap is then compared against the fixed one. -Reported per method: the mean, 5th percentile, median, 95th percentile, -minimum and maximum of the per-class Dice scores, the number of mislabeled -voxels, and the wall-clock registration time. The unregistered moving image, -resampled onto the fixed grid and segmented the same way, supplies the "before -registration" reference row. - -Note that segmenting each registered image separately means the scores include -segmentation variability on the warped volumes, not the geometric error of the -transform alone. It also costs one GPU segmentation per method. +Accuracy is measured two ways. The primary metric is target registration +error: DIR-Lab ships 300 expert landmarks for the extreme phases (T00 and T50) +of every case, so each fixed-image landmark is mapped through the registration +transform and compared, in millimeters, against its moving-image counterpart. +The secondary metric is label overlap: ``SegmentNVSegmentCTMRI`` segments the +fixed and moving images once each, and the moving labelmap is warped onto the +fixed grid by every transform, so the Dice scores reflect the transform rather +than segmentation variability on re-segmented warped volumes. The moving image +and labelmap resampled onto the fixed grid without registration supply the +"before registration" reference row for both metrics. + +Reported per method: the mean, standard deviation, 95th percentile and maximum +landmark error in millimeters; the mean, 5th percentile, median, 95th +percentile, minimum and maximum of the per-class Dice scores; the number of +mislabeled voxels; and the wall-clock registration time. Finetuning artifacts (dataset JSON, YAML config, checkpoint tree) are written under ``tutorials/network_weights/icon_dirlab_4dct``. The final checkpoint is ``tutorials/network_weights/icon_dirlab_4dct/icon_dirlab_4dct_model/checkpoints/ -Finetune_multi_final.trch``, the path returned by +network_weights_final.trch``, the path returned by ``WorkflowFinetuneICONRegistration.expected_weights_path()``. That directory is deleted at the start of every run, so each run finetunes from scratch; see the comment above the ``shutil.rmtree`` call for how to reuse a previous run. @@ -34,7 +37,8 @@ Data Required ------------- Full data: ``data/DirLab-4DCT`` (all 10 cases, converted to HU ``.mha`` by -``data/DirLab-4DCT/fix_downloaded_data.py``) +``data/DirLab-4DCT/fix_downloaded_data.py``), including the raw +``downloaded_data/Case1Pack/ExtremePhases`` landmark files Test data: ``data/test/DirLab-4DCT`` """ @@ -55,7 +59,7 @@ PhysioTwin4DBase, RegisterImagesBase, RegisterImagesGreedy, - RegisterImagesGreedyICON, + RegisterImagesICON, SegmentNVSegmentCTMRI, TestTools, TransformTools, @@ -82,14 +86,20 @@ weights_dir = tutorials_dir / "network_weights" finetune_name = "icon_dirlab_4dct" + run_finetuning = True + test_mode = TestTools.running_as_test() if test_mode: data_dir = repo_root / "data" / "test" / "DirLab-4DCT" number_of_iterations_greedy: Optional[list[int]] = [1, 0] - epochs = 5 + epochs = 1 else: data_dir = repo_root / "data" / "DirLab-4DCT" - number_of_iterations_greedy = None # Greedy defaults + number_of_iterations_greedy = [60, 30, 20] # Greedy defaults + # 90 training frames at batch_size 4 is 22 optimizer steps per epoch, so + # 100 epochs is ~2200 steps at a 5e-5 learning rate. Far fewer than + # that leaves the finetuned weights statistically indistinguishable + # from the stock weights they started from. epochs = 100 log_level = logging.INFO @@ -97,13 +107,22 @@ output_dir.mkdir(parents=True, exist_ok=True) - # Held-out evaluation pair (Case 1 is excluded from finetuning) - fixed_file = data_dir / "Case1Pack_T70.mha" - moving_file = data_dir / "Case1Pack_T30.mha" - missing = [str(p) for p in (fixed_file, moving_file) if not p.exists()] + # Held-out evaluation pair (Case 1 is excluded from finetuning). T00 and + # T50 are the extreme inhale/exhale phases, the only pair DIR-Lab supplies + # expert landmarks for. + fixed_file = data_dir / "Case1Pack_T50.mha" + moving_file = data_dir / "Case1Pack_T00.mha" + landmark_dir = data_dir / "downloaded_data" / "Case1Pack" / "ExtremePhases" + fixed_landmark_file = landmark_dir / "Case1_300_T50_xyz.txt" + moving_landmark_file = landmark_dir / "Case1_300_T00_xyz.txt" + missing = [ + str(p) + for p in (fixed_file, moving_file, fixed_landmark_file, moving_landmark_file) + if not p.exists() + ] if missing: raise FileNotFoundError( - f"Missing DirLab phase images: {missing}.\n" + f"Missing DirLab phase images or landmarks: {missing}.\n" "See data/DirLab-4DCT/README.md for download instructions." ) @@ -140,35 +159,84 @@ # if not weights_path.exists(): # weights_path = workflow.process() experiment_dir = weights_dir / finetune_name - if experiment_dir.exists(): - reporter.log_info("Removing previous finetuning outputs: %s", experiment_dir) - shutil.rmtree(experiment_dir) - - # DIR-Lab ships no segmentations, so no labelmaps or masks are supplied and - # the Dice loss must be disabled: uniGradICON requires a ``segmentation`` - # field on every dataset entry when dice_loss_weight > 0. - workflow = WorkflowFinetuneICONRegistration( - subject_image_files=list(subject_image_files.values()), - output_dir=weights_dir, - finetune_name=finetune_name, - subject_ids=list(subject_image_files.keys()), - epochs=epochs, - dice_loss_weight=0.0, - log_level=log_level, - ) - weights_path = workflow.process() + if run_finetuning: + if experiment_dir.exists(): + reporter.log_info( + "Removing previous finetuning outputs: %s", experiment_dir + ) + shutil.rmtree(experiment_dir) + + # DIR-Lab ships no segmentations, so no labelmaps or masks are supplied and + # the Dice loss must be disabled: uniGradICON requires a ``segmentation`` + # field on every dataset entry when dice_loss_weight > 0. + # + # lncc_sigma matches the sigma RegisterImagesICON uses at inference, so + # finetuning optimizes the similarity this comparison scores. + workflow = WorkflowFinetuneICONRegistration( + subject_image_files=list(subject_image_files.values()), + output_dir=weights_dir, + finetune_name=finetune_name, + subject_ids=list(subject_image_files.keys()), + epochs=epochs, + dice_loss_weight=0.0, + lncc_sigma=5, + log_level=log_level, + ) + weights_path = workflow.process() + else: + weights_path = ( + Path(__file__).resolve().parent + / "network_weights/icon_dirlab_4dct/icon_dirlab_4dct_model/checkpoints/network_weights_final.trch" + ) # Registration comparison fixed_image = itk.imread(str(fixed_file), pixel_type=itk.F) moving_image = itk.imread(str(moving_file), pixel_type=itk.F) transform_tools = TransformTools() - # Every image is segmented independently: the fixed image once, then each - # registered image after warping. The Dice scores therefore include - # whatever the segmenter does differently on each warped volume, not only - # the geometric error of the transform. + def read_landmarks(landmark_file: Path, image: itk.Image) -> np.ndarray: + """Read a DIR-Lab landmark file as an (N, 3) array of world points. + + Each line holds one 1-based voxel index as ``x y z``. + """ + indices = np.loadtxt(landmark_file, dtype=int) - 1 + return np.array( + [ + image.TransformIndexToPhysicalPoint([int(v) for v in index]) + for index in indices + ] + ) + + fixed_landmarks = read_landmarks(fixed_landmark_file, fixed_image) + moving_landmarks = read_landmarks(moving_landmark_file, moving_image) + + def landmark_metrics(errors_mm: np.ndarray) -> dict[str, Any]: + """Summarize per-landmark target registration errors, in millimeters.""" + return { + "tre_mean": float(errors_mm.mean()), + "tre_std": float(errors_mm.std()), + "tre_p95": float(np.percentile(errors_mm, 95)), + "tre_max": float(errors_mm.max()), + } + + def landmark_errors(transform: itk.Transform) -> np.ndarray: + """Distance from each mapped fixed landmark to its moving counterpart. + + ``forward_transform`` is the resampling transform: it maps points on the + fixed grid back into moving space, which is the direction the landmark + correspondences are defined in. + """ + mapped = np.array( + [transform.TransformPoint(tuple(point)) for point in fixed_landmarks] + ) + return np.asarray(np.linalg.norm(mapped - moving_landmarks, axis=1)) + + # Each image is segmented once and the moving labelmap is warped by every + # transform, so Dice reflects the transform rather than what the segmenter + # does differently on each interpolated volume. segmenter = SegmentNVSegmentCTMRI(log_level=log_level) fixed_labelmap = segmenter.segment(fixed_image)["labelmap"] + moving_labelmap = segmenter.segment(moving_image)["labelmap"] fixed_labels = itk.array_from_image(fixed_labelmap) def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: @@ -199,24 +267,32 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: "mislabeled_voxels": int(np.count_nonzero(fixed_labels != labels)), } - # Reference row: the moving image on the fixed grid, unregistered. + # Reference row: the moving image and its labelmap on the fixed grid, + # unregistered. unregistered_image = itk.resample_image_filter( moving_image, ReferenceImage=fixed_image, UseReferenceImage=True, ) + unregistered_labelmap = itk.resample_image_filter( + moving_labelmap, + Interpolator=itk.NearestNeighborInterpolateImageFunction.New(moving_labelmap), + ReferenceImage=fixed_image, + UseReferenceImage=True, + ) - registered_images: dict[str, itk.Image] = {} - labelmaps: dict[str, itk.Image] = { - "unregistered": segmenter.segment(unregistered_image)["labelmap"] - } + registered_images: dict[str, itk.Image] = {"unregistered": unregistered_image} + labelmaps: dict[str, itk.Image] = {"unregistered": unregistered_labelmap} rows: list[dict[str, Any]] = [ { "method": "unregistered", "weights": "-", "registration_time_s": None, "loss": None, - **overlap_metrics(labelmaps["unregistered"]), + **landmark_metrics( + np.linalg.norm(fixed_landmarks - moving_landmarks, axis=1) + ), + **overlap_metrics(unregistered_labelmap), } ] for method_name, method_weights in ( @@ -227,19 +303,18 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: registrar: RegisterImagesBase if method_name == "greedy": registrar = RegisterImagesGreedy(log_level=log_level) + registrar.set_transform_type("Deformable") if number_of_iterations_greedy is not None: registrar.set_number_of_iterations(number_of_iterations_greedy) else: - registrar = RegisterImagesGreedyICON(log_level=log_level) - if number_of_iterations_greedy is not None: - registrar.greedy.set_number_of_iterations(number_of_iterations_greedy) + registrar = RegisterImagesICON(log_level=log_level) # None, not 0: icon_registration rejects 0 and takes None to mean # "no test-time finetuning steps", so the comparison reflects what # each set of weights predicts rather than per-pair optimization. - registrar.icon.set_number_of_iterations(None) - registrar.icon.set_mass_preservation(True) # For non-contrast CT + registrar.set_number_of_iterations(None) + registrar.set_mass_preservation(True) # For non-contrast CT if method_weights is not None: - registrar.icon.set_weights_path(str(method_weights)) + registrar.set_weights_path(str(method_weights)) registrar.set_modality("ct") registrar.set_fixed_image(fixed_image) @@ -247,17 +322,22 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: result = registrar.register(moving_image) elapsed_s = time.perf_counter() - start_time - registered = transform_tools.transform_image( + registered_images[method_name] = transform_tools.transform_image( moving_image, result["forward_transform"], fixed_image ) - registered_images[method_name] = registered - labelmaps[method_name] = segmenter.segment(registered)["labelmap"] + labelmaps[method_name] = transform_tools.transform_image( + moving_labelmap, + result["forward_transform"], + fixed_image, + interpolation_method="nearest", + ) rows.append( { "method": method_name, "weights": str(method_weights) if method_weights else "-", "registration_time_s": elapsed_s, "loss": float(result["loss"]), + **landmark_metrics(landmark_errors(result["forward_transform"])), **overlap_metrics(labelmaps[method_name]), } ) @@ -287,10 +367,27 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: # Reporting reporter.log_info( - "Case1Pack_T30 -> Case1Pack_T70, per-class Dice against the fixed labelmap" + "Case1Pack_T00 -> Case1Pack_T50, error at %d expert landmarks, mm", + len(fixed_landmarks), + ) + reporter.log_info( + " %-13s %7s %7s %7s %7s %9s", "method", "mean", "std", "p95", "max", "time_s" ) + for row in rows: + elapsed = row["registration_time_s"] + reporter.log_info( + " %-13s %7.2f %7.2f %7.2f %7.2f %9s", + row["method"], + row["tre_mean"], + row["tre_std"], + row["tre_p95"], + row["tre_max"], + "-" if elapsed is None else f"{float(elapsed):.1f}", + ) + + reporter.log_info("Per-class Dice of the warped moving labelmap against the fixed") reporter.log_info( - " %-13s %7s %7s %7s %7s %7s %7s %7s %12s %9s", + " %-13s %7s %7s %7s %7s %7s %7s %7s %12s", "method", "classes", "mean", @@ -300,12 +397,10 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: "min", "max", "mislabeled", - "time_s", ) for row in rows: - elapsed = row["registration_time_s"] reporter.log_info( - " %-13s %7d %7.4f %7.4f %7.4f %7.4f %7.4f %7.4f %12d %9s", + " %-13s %7d %7.4f %7.4f %7.4f %7.4f %7.4f %7.4f %12d", row["method"], row["n_classes"], row["dice_mean"], @@ -315,7 +410,6 @@ def overlap_metrics(labelmap: itk.Image) -> dict[str, Any]: row["dice_min"], row["dice_max"], row["mislabeled_voxels"], - "-" if elapsed is None else f"{float(elapsed):.1f}", ) reporter.log_info("Wrote summary: %s", summary_file) diff --git a/tutorials/tutorial_03_lung_reconstruct_highres_4d_ct.py b/tutorials/tutorial_03_lung_reconstruct_highres_4d_ct.py index a7dfbe0..73378dd 100644 --- a/tutorials/tutorial_03_lung_reconstruct_highres_4d_ct.py +++ b/tutorials/tutorial_03_lung_reconstruct_highres_4d_ct.py @@ -13,7 +13,7 @@ Full data: ``data/DirLab-4DCT/Case1Pack_T??.mha`` Test data: ``data/test/DirLab-4DCT/Case1Pack_T??.mha`` ICON weights: Tutorial 2 output -(``tutorials/network_weights/icon_dirlab_4dct/.../Finetune_multi_final.trch``), +(``tutorials/network_weights/icon_dirlab_4dct/.../network_weights_final.trch``), optional — the stock uniGradICON weights are used when it is absent. Outputs (under ``tutorials/output/tutorial_03_lung/``) @@ -66,7 +66,7 @@ / "icon_dirlab_4dct" / "icon_dirlab_4dct_model" / "checkpoints" - / "Finetune_multi_final.trch" + / "network_weights_final.trch" ) test_mode = TestTools.running_as_test() diff --git a/tutorials/tutorial_04_heart_ct_to_vtk.py b/tutorials/tutorial_04_heart_ct_to_vtk.py index 7c3cd29..069f0d9 100644 --- a/tutorials/tutorial_04_heart_ct_to_vtk.py +++ b/tutorials/tutorial_04_heart_ct_to_vtk.py @@ -24,7 +24,9 @@ from physiotwin4d import ( ContourTools, + SegmentAnatomyBase, SegmentChestTotalSegmentatorWithContrast, + SegmentHeartSimpleware, TestTools, WorkflowConvertImageToVTK, ) @@ -51,6 +53,9 @@ save_group_surfaces = True save_label_surfaces = True + use_simpleware = False + use_totalsegmentator_academic_license = True + test_mode = TestTools.running_as_test() if test_mode: data_dir = repo_root / "data" / "test" / "slicer_heart_small" @@ -61,8 +66,18 @@ log_level = logging.INFO - segmentation_method = SegmentChestTotalSegmentatorWithContrast(log_level=log_level) - segmentation_method.set_has_academic_license(True) + if use_simpleware: + segmentation_method: SegmentAnatomyBase = SegmentHeartSimpleware( + log_level=log_level + ) + else: + total_segmentation_method = SegmentChestTotalSegmentatorWithContrast( + log_level=log_level + ) + total_segmentation_method.set_has_academic_license( + use_totalsegmentator_academic_license + ) + segmentation_method = total_segmentation_method # Directory setup and data reading diff --git a/tutorials/tutorial_08_lung_fit_model_to_4d_patients.py b/tutorials/tutorial_08_lung_fit_model_to_4d_patients.py index 10fb2a9..2f6f9a2 100644 --- a/tutorials/tutorial_08_lung_fit_model_to_4d_patients.py +++ b/tutorials/tutorial_08_lung_fit_model_to_4d_patients.py @@ -32,7 +32,7 @@ ``pca_mean_surface.vtp``) ICON weights: Tutorial 2 output (``network_weights/icon_dirlab_4dct/icon_dirlab_4dct_model/checkpoints/ -Finetune_multi_final.trch``), optional — the stock uniGradICON weights are used +network_weights_final.trch``), optional — the stock uniGradICON weights are used when it is absent. Outputs (per case, under ``output/tutorial_08_lung//``) @@ -98,7 +98,7 @@ / "icon_dirlab_4dct" / "icon_dirlab_4dct_model" / "checkpoints" - / "Finetune_multi_final.trch" + / "network_weights_final.trch" ) # Phase the SSM is fitted to; Tutorial 6 builds the lung PCA model from the diff --git a/utils/patches/unigradicon_augment_identity.patch b/utils/patches/unigradicon_augment_identity.patch new file mode 100644 index 0000000..04ecdf9 --- /dev/null +++ b/utils/patches/unigradicon_augment_identity.patch @@ -0,0 +1,74 @@ +Subject: [PATCH] finetuning: stop randomly permuting and flipping axes in augment() + +`augment()` builds its "identity" affine by drawing a random axis permutation +and then a random sign per axis, giving 48 orientations (6 permutations x 8 sign +patterns) of which only one matches the data. Image A and image B share the +permutation, so a training pair stays internally consistent, but inference never +reorients anything -- `unigradicon.preprocess` and `icon_registration`'s +`register_pair` both work on the array as read, with no direction handling. The +permuted and mirrored poses therefore never occur at test time. + +For a foundation-model pretraining run that is arguably useful regularization. +For finetuning it is not: a finetune is a few hundred optimizer steps at a low +learning rate on a task-specific cohort, and spending most of them on +orientations the model will never be asked about leaves the checkpoint +statistically indistinguishable from the starting weights. Observed on a +DIR-Lab 4D CT lung finetune (9 cases, 90 volumes, 220 steps at 5e-5): mean Dice +on a held-out case moved by 4e-4 versus stock uniGradICON. + +This patch keeps the random affine augmentation -- the `0.05 * randn` noise, and +the independent `noise_A` / `noise_B` draws that make the pair differ in detail +-- and only replaces the randomized orientation with a true identity. + +Applies to the `feat-add-finetuning` branch of +https://github.com/uncbiag/uniGradICON. + +Note for this repository: `pyproject.toml` installs uniGradICON from that branch +via `[tool.uv.sources]`, so any reinstall reverts a locally applied copy of this +patch until it lands upstream. Re-apply with: + + py -c "import unigradicon, os; print(os.path.dirname(os.path.dirname(unigradicon.__file__)))" + cd + git apply -p2 /utils/patches/unigradicon_augment_identity.patch + +`-p2` strips the leading ``a/src/``, leaving ``unigradicon/finetuning/ +finetune.py``, which is the installed layout. ``git apply -R -p2`` reverts it, +and either direction fails cleanly if the file has already been changed. + +--- a/src/unigradicon/finetuning/finetune.py ++++ b/src/unigradicon/finetuning/finetune.py +@@ -53,23 +53,21 @@ + Images are warped with bilinear interpolation; segmentations and masks + use nearest interpolation to preserve label values. + +- Both images share the same random flip/permutation but have slightly +- different affine noise, so they share orientation but differ in detail. ++ The two images of a pair get independent affine noise around a common ++ identity, so they stay in the same orientation but differ in detail. ++ ++ The augmentation deliberately does not permute or flip axes. Inference ++ applies no reorientation, so the 48 permuted/mirrored orientations never ++ occur at test time and training on them spends the finetuning budget on ++ poses the model is never asked about. + """ + device = batch[PairKeys.IMAGE_A].device + batch_size = batch[PairKeys.IMAGE_A].shape[0] + +- identity_list = [] +- for _ in range(batch_size): +- identity = torch.zeros((1, 3, 4), dtype=torch.float32, device=device) +- idxs = {0, 1, 2} +- for j in range(3): +- k = random.choice(list(idxs)) +- idxs.remove(k) +- identity[0, j, k] = 1 +- identity = identity * (torch.randint_like(identity, 0, 2) * 2 - 1) +- identity_list.append(identity) +- identity = torch.cat(identity_list) ++ identity = torch.zeros((batch_size, 3, 4), dtype=torch.float32, device=device) ++ identity[:, 0, 0] = 1 ++ identity[:, 1, 1] = 1 ++ identity[:, 2, 2] = 1 + + noise_A = torch.randn((batch_size, 3, 4), device=device) + forward_A = identity + 0.05 * noise_A