Dev

.gitignore

@@ -161,6 +161,10 @@ models/
 cv_summary.csv
 patch_config.json
 
+# Patch-workflow run byproducts (regenerated by the notebooks; e.g. research/*)
+preprocessed/
+cv_results_*/
+
 # Claude - share skills/settings, keep local config private
 .claude/settings.local.json
 .claude/*.local.*

fastMONAI/__init__.py

@@ -1 +1 @@
-__version__ = "0.9.1"
+__version__ = "0.9.2"

fastMONAI/_modidx.py

@@ -125,6 +125,7 @@
                                                                                  'fastMONAI/utils.py'),
                                  'fastMONAI.utils._to_jsonable': ('utils.html#_to_jsonable', 'fastMONAI/utils.py'),
                                  'fastMONAI.utils.create_mlflow_callback': ('utils.html#create_mlflow_callback', 'fastMONAI/utils.py'),
+                                 'fastMONAI.utils.find_fold_learners': ('utils.html#find_fold_learners', 'fastMONAI/utils.py'),
                                  'fastMONAI.utils.load_patch_variables': ('utils.html#load_patch_variables', 'fastMONAI/utils.py'),
                                  'fastMONAI.utils.load_variables': ('utils.html#load_variables', 'fastMONAI/utils.py'),
                                  'fastMONAI.utils.print_colab_gpu_info': ('utils.html#print_colab_gpu_info', 'fastMONAI/utils.py'),
@@ -418,6 +419,12 @@
                                                                                         'fastMONAI/vision_loss.py'),
                                        'fastMONAI.vision_loss.CustomLoss.decodes': ( 'vision_loss_functions.html#customloss.decodes',
                                                                                      'fastMONAI/vision_loss.py'),
+                                       'fastMONAI.vision_loss.DynUNetDSAdapter': ( 'vision_loss_functions.html#dynunetdsadapter',
+                                                                                   'fastMONAI/vision_loss.py'),
+                                       'fastMONAI.vision_loss.DynUNetDSAdapter.__init__': ( 'vision_loss_functions.html#dynunetdsadapter.__init__',
+                                                                                            'fastMONAI/vision_loss.py'),
+                                       'fastMONAI.vision_loss.DynUNetDSAdapter.forward': ( 'vision_loss_functions.html#dynunetdsadapter.forward',
+                                                                                           'fastMONAI/vision_loss.py'),
                                        'fastMONAI.vision_loss.TverskyFocalLoss': ( 'vision_loss_functions.html#tverskyfocalloss',
                                                                                    'fastMONAI/vision_loss.py'),
                                        'fastMONAI.vision_loss.TverskyFocalLoss.__init__': ( 'vision_loss_functions.html#tverskyfocalloss.__init__',
@@ -464,6 +471,8 @@
                                                                                            'fastMONAI/vision_metrics.py'),
                                           'fastMONAI.vision_metrics.binary_signed_rve': ( 'vision_metrics.html#binary_signed_rve',
                                                                                           'fastMONAI/vision_metrics.py'),
+                                          'fastMONAI.vision_metrics.calculate_all_metrics': ( 'vision_metrics.html#calculate_all_metrics',
+                                                                                              'fastMONAI/vision_metrics.py'),
                                           'fastMONAI.vision_metrics.calculate_confusion_metrics': ( 'vision_metrics.html#calculate_confusion_metrics',
                                                                                                     'fastMONAI/vision_metrics.py'),
                                           'fastMONAI.vision_metrics.calculate_dsc': ( 'vision_metrics.html#calculate_dsc',
@@ -478,6 +487,8 @@
                                                                                              'fastMONAI/vision_metrics.py'),
                                           'fastMONAI.vision_metrics.calculate_surface_metrics': ( 'vision_metrics.html#calculate_surface_metrics',
                                                                                                   'fastMONAI/vision_metrics.py'),
+                                          'fastMONAI.vision_metrics.evaluate_segmentations': ( 'vision_metrics.html#evaluate_segmentations',
+                                                                                               'fastMONAI/vision_metrics.py'),
                                           'fastMONAI.vision_metrics.multi_dice_score': ( 'vision_metrics.html#multi_dice_score',
                                                                                          'fastMONAI/vision_metrics.py'),
                                           'fastMONAI.vision_metrics.multi_lesion_detection_rate': ( 'vision_metrics.html#multi_lesion_detection_rate',

fastMONAI/utils.py

@@ -2,7 +2,8 @@
 
 # %% auto #0
 __all__ = ['set_mlflow_tracking_uri', 'store_variables', 'load_variables', 'store_patch_variables', 'load_patch_variables',
-           'print_colab_gpu_info', 'ModelTrackingCallback', 'create_mlflow_callback', 'MLflowUIManager']
+           'print_colab_gpu_info', 'ModelTrackingCallback', 'create_mlflow_callback', 'find_fold_learners',
+           'MLflowUIManager']
 
 # %% ../nbs/07_utils.ipynb #ac941446-cf7e-4f5c-ace0-a36fb078022f
 import pickle
@@ -731,10 +732,13 @@ def log_metrics_table(self, df, metrics=None, key='validation_metrics', display=
 
         summary_data = []
         for metric in metrics:
-            mean = df[metric].mean()
-            std = df[metric].std()
+            # Drop +/-inf (e.g. surface distances on one_empty cases) so they don't poison
+            # the summary; pandas already skips NaN. Show 'N/A' if nothing finite remains.
+            col = df[metric].replace([float('inf'), float('-inf')], float('nan'))
+            mean, std = col.mean(), col.std()
+            mean_str = f'{mean:.4f}' if not math.isnan(mean) else 'N/A'
             std_str = f'{std:.4f}' if not math.isnan(std) else 'N/A'
-            summary_data.append([metric, f'{mean:.4f}', std_str])
+            summary_data.append([metric, mean_str, std_str])
 
         fig, ax = plt.subplots(figsize=(5, len(metrics) * 0.8 + 0.5))
         ax.axis('off')
@@ -881,6 +885,85 @@ def create_mlflow_callback(
         log_split=log_split,
     )
 
+# %% ../nbs/07_utils.ipynb #6e7ed65b
+def find_fold_learners(experiment_name, artifact_path="model/best_learner.pkl",
+                       n_folds=None, fold_tag="fold"):
+    """Discover one trained learner checkpoint per cross-validation fold from an MLflow experiment.
+
+    Reads back what ``ModelTrackingCallback`` logs per run: takes the most recent run for each
+    fold in ``experiment_name`` and downloads ``artifact_path`` from it, returning
+    ``{fold: local_path}`` sorted by fold. Folds are identified by each run's ``tags.<fold_tag>``,
+    falling back to a ``<fold_tag>_<n>`` run name. Warns if the selected folds span more than one
+    ``dataset_version`` tag (a sign of a partial retrain).
+
+    Pair it with the soft-vote ensembling in ``patch_inference`` -- but the returned values are
+    PATHS, so load them first::
+
+        from fastai.learner import load_learner
+        paths = find_fold_learners("vs5f_unet")
+        learners = [load_learner(p) for p in paths.values()]
+        preds = patch_inference(learner=learners, config=config, file_paths=cases)
+
+    Args:
+        experiment_name: MLflow experiment to search (e.g. ``"vs5f_unet"``).
+        artifact_path: run-relative artifact to download per fold (default
+            ``"model/best_learner.pkl"``, what ``ModelTrackingCallback`` logs).
+        n_folds: optional expected folds for a missing-fold warning -- an int ``n`` (labels
+            ``1..n``) or an explicit iterable of labels (use this for 0-indexed folds). ``None``
+            makes no assumption and just returns what was found.
+        fold_tag: run tag naming the fold; also the ``<fold_tag>_`` run-name prefix fallback.
+
+    Returns:
+        ``{fold_int: local_path}`` sorted by fold, or ``{}`` if the experiment is missing or
+        MLflow is unavailable.
+    """
+    from mlflow.exceptions import MlflowException
+    try:
+        _ensure_fastmonai_tracking_uri()  # respect a user-configured tracking server
+        exp = mlflow.get_experiment_by_name(experiment_name)
+        if exp is None:
+            print(f"No MLflow experiment {experiment_name!r}. Nothing to load.")
+            return {}
+        runs = mlflow.search_runs(experiment_ids=[exp.experiment_id],
+                                  order_by=["attributes.start_time DESC"])
+    except (ImportError, MlflowException) as e:
+        print(f"MLflow lookup for {experiment_name!r} skipped: {e}")
+        return {}
+
+    out, fold_dsv = {}, {}
+    prefix = f"{fold_tag}_"
+    for _, row in runs.iterrows():
+        fold = row.get(f"tags.{fold_tag}")
+        # pandas yields NaN (a float, != itself), not None, when the column exists but this
+        # row lacks the tag -- fall back to the run name instead of dropping the run.
+        if fold is None or (isinstance(fold, float) and fold != fold):
+            name = row.get("tags.mlflow.runName") or ""
+            fold = name[len(prefix):] if name.startswith(prefix) else None
+        try:
+            fold = int(fold)
+        except (TypeError, ValueError):
+            continue
+        if fold in out:
+            continue  # newest run already kept for this fold (runs are start_time DESC)
+        try:
+            out[fold] = mlflow.artifacts.download_artifacts(
+                run_id=row["run_id"], artifact_path=artifact_path)
+        except Exception:
+            continue  # fall through to the next-newest run for this fold
+        fold_dsv[fold] = row.get("tags.dataset_version")
+
+    print(f"Found fold learners for {sorted(out)} in {experiment_name!r}.")
+    if n_folds is not None:
+        expected = range(1, n_folds + 1) if isinstance(n_folds, int) else n_folds
+        missing = [f for f in expected if f not in out]
+        if missing:
+            print(f"WARNING: no checkpoint for folds {missing}; ensemble will use {len(out)} model(s).")
+    distinct_dsv = {v for v in fold_dsv.values() if isinstance(v, str) and v}
+    if len(distinct_dsv) > 1:
+        shown = {f: v[:8] for f, v in sorted(fold_dsv.items()) if isinstance(v, str) and v}
+        print(f"WARNING: selected folds span multiple dataset_versions {shown} - possible partial retrain.")
+    return dict(sorted(out.items()))
+
 # %% ../nbs/07_utils.ipynb #a3d2d585
 import subprocess
 import threading

fastMONAI/vision_loss.py

@@ -1,7 +1,7 @@
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/04_vision_loss_functions.ipynb.
 
 # %% auto #0
-__all__ = ['CustomLoss', 'TverskyFocalLoss']
+__all__ = ['CustomLoss', 'TverskyFocalLoss', 'DynUNetDSAdapter']
 
 # %% ../nbs/04_vision_loss_functions.ipynb #e8eee6b3-75c3-4468-8645-8f6a42a76bfe
 import torch
@@ -105,3 +105,31 @@ def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         total_loss: torch.Tensor = tversky_loss ** self.gamma
 
         return total_loss
+
+# %% ../nbs/04_vision_loss_functions.ipynb #acf08940
+class DynUNetDSAdapter(torch.nn.Module):
+    """Wrap a MONAI ``DynUNet`` with deep supervision so its output matches what
+    ``monai.losses.DeepSupervisionLoss`` (used via ``CustomLoss``) expects.
+
+    A ``DynUNet(deep_supervision=True)`` returns a *stacked* ``[B, N, C, ...]`` tensor in
+    training (a TorchScript workaround), but ``DeepSupervisionLoss`` only unbinds a *list* of
+    tensors -- a stacked tensor would fall through to the base loss and be scored incorrectly.
+    This adapter unbinds the stack into that list in training, and passes the plain prediction
+    through in eval mode (MONAI's ``DynUNet`` docstring prescribes exactly this ``torch.unbind``).
+
+    Note: wrapping registers the network as ``self.model``, so it adds a ``model.`` prefix to
+    state-dict keys -- relevant when loading raw ``.pth`` checkpoints of a wrapped model.
+    """
+
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+
+    def forward(self, x):
+        out = self.model(x)
+        # Unbind only in training, and only when the wrapped net actually stacks its
+        # deep-supervision heads. Gating on `deep_supervision` (not output rank) keeps this
+        # correct for 2D and 3D alike, and fails loudly if wrapped around a non-DynUNet in training.
+        if self.model.training and self.model.deep_supervision:
+            return list(torch.unbind(out, dim=1))
+        return out

fastMONAI/vision_metrics.py

@@ -5,14 +5,17 @@
            'calculate_confusion_metrics', 'binary_sensitivity', 'multi_sensitivity', 'binary_precision',
            'multi_precision', 'calculate_lesion_detection_rate', 'binary_lesion_detection_rate',
            'multi_lesion_detection_rate', 'calculate_signed_rve', 'binary_signed_rve', 'multi_signed_rve',
-           'calculate_surface_metrics', 'AccumulatedDice', 'AccumulatedMultiDice', 'EMACheckpoint']
+           'calculate_surface_metrics', 'calculate_all_metrics', 'evaluate_segmentations', 'AccumulatedDice',
+           'AccumulatedMultiDice', 'EMACheckpoint']
 
 # %% ../nbs/05_vision_metrics.ipynb #8b6a83ac
 import warnings
 import torch
 import numpy as np
+import pandas as pd
 from monai.metrics import compute_hausdorff_distance, compute_dice, get_confusion_matrix, compute_confusion_matrix_metric, compute_average_surface_distance, compute_surface_dice
 from scipy.ndimage import label as scipy_label
+from pathlib import Path
 from .vision_data import pred_to_binary_mask, batch_pred_to_multiclass_mask
 from fastai.learner import Metric
 from fastai.callback.tracker import TrackerCallback
@@ -506,6 +509,97 @@ def calculate_surface_metrics(pred, targ, spacing_mm,
             **nsd,
             "status": "ok", **base}
 
+# %% ../nbs/05_vision_metrics.ipynb #ed336bbf
+def calculate_all_metrics(pred, targ, spacing_mm, *, nsd_tolerances_mm=(0.5, 1.0, 2.0),
+                          binary: bool = True) -> dict:
+    """Full per-case metric panel for one segmentation vs its ground truth.
+
+    Orchestrates the individual ``calculate_*`` functions into one dict: Dice, sensitivity,
+    precision, lesion detection rate, signed RVE, and the spacing-aware surface metrics
+    (ASSD / HD95 / NSD in mm). For a whole dataset, use ``evaluate_segmentations``.
+
+    Args:
+        pred: predicted mask (binary / argmax label). torch.Tensor or np.ndarray, shaped
+            [W,H,D], [C,W,H,D] or [B,C,W,H,D] with singleton batch/channel dims.
+        targ: ground-truth mask, same shape convention, on the SAME voxel grid as ``pred``.
+        spacing_mm: voxel spacing in mm in array-axis (i,j,k) order matching ``targ``. Pass it
+            explicitly (read from the GT file); a permuted spacing silently corrupts the
+            surface distances.
+        nsd_tolerances_mm: NSD tolerance(s) in mm, forwarded to ``calculate_surface_metrics``.
+        binary: single-foreground-class segmentation (the only mode supported). The confusion
+            metrics run on the single label channel -- MONAI keeps a size-1 channel when
+            ``include_background=False``, so a single-channel input is scored correctly.
+
+    Returns:
+        dict: ``dsc``, ``sensitivity``, ``precision``, ``ldr``, ``rve``, the ``nsd_tau*_mm``
+        values, ``assd_mm``, ``hd95_mm``, ``surface_status`` and ``spacing_mm``.
+    """
+    if not binary:
+        raise NotImplementedError("calculate_all_metrics supports binary (single foreground "
+                                  "class) segmentation only.")
+
+    pred_t = pred if isinstance(pred, torch.Tensor) else torch.as_tensor(np.asarray(pred))
+    targ_t = targ if isinstance(targ, torch.Tensor) else torch.as_tensor(np.asarray(targ))
+    while pred_t.ndim < 5: pred_t = pred_t.unsqueeze(0)   # -> [B,C,W,H,D]
+    while targ_t.ndim < 5: targ_t = targ_t.unsqueeze(0)
+    if pred_t.shape[-3:] != targ_t.shape[-3:]:
+        raise ValueError("pred and targ must share a voxel grid; got spatial shapes "
+                         f"{tuple(pred_t.shape[-3:])} vs {tuple(targ_t.shape[-3:])}.")
+    pred_t, targ_t = pred_t.float(), targ_t.float()
+
+    sm = calculate_surface_metrics(pred_t, targ_t, spacing_mm, nsd_tolerances_mm=nsd_tolerances_mm)
+    row = {
+        "dsc": calculate_dsc(pred_t, targ_t).nanmean().item(),
+        "sensitivity": calculate_confusion_metrics(pred_t, targ_t, "sensitivity").nanmean().item(),
+        "precision": calculate_confusion_metrics(pred_t, targ_t, "precision").nanmean().item(),
+        "ldr": calculate_lesion_detection_rate(pred_t, targ_t).nanmean().item(),
+        "rve": calculate_signed_rve(pred_t, targ_t).nanmean().item(),
+    }
+    row.update({k: sm[k] for k in sm if k.startswith("nsd_tau")})
+    row["assd_mm"], row["hd95_mm"] = sm["assd_mm"], sm["hd95_mm"]
+    row["surface_status"], row["spacing_mm"] = sm["status"], sm["spacing_mm"]
+    return row
+
+# %% ../nbs/05_vision_metrics.ipynb #2e526be3
+def evaluate_segmentations(preds, gt_paths, case_ids=None, *,
+                           nsd_tolerances_mm=(0.5, 1.0, 2.0), binary: bool = True):
+    """Score a batch of predictions against ground-truth masks into a per-case DataFrame.
+
+    For each (prediction, ground-truth) pair, the GT mask AND its voxel spacing are read from a
+    SINGLE object, so the array and its spacing can never disagree -- reading spacing from a
+    separately-loaded file is a silent source of axis-permutation errors that corrupt the surface
+    metrics. Predictions may be tensors (as returned by ``patch_inference``) or paths to saved
+    masks, so predictions written with ``patch_inference(save_dir=...)`` can be re-scored without
+    re-running inference.
+
+    Args:
+        preds: sequence of predicted masks -- torch.Tensor / np.ndarray, or paths to mask files.
+        gt_paths: sequence of ground-truth mask file paths (one per prediction, same voxel grid).
+        case_ids: optional row labels; default is each GT file's name stem.
+        nsd_tolerances_mm: NSD tolerance(s) in mm, forwarded to the surface metrics.
+        binary: binary (single foreground class) segmentation only.
+
+    Returns:
+        pandas.DataFrame with one row per case: ``case_id`` plus the columns from
+        ``calculate_all_metrics``.
+    """
+    import torchio as tio  # lazy: only this benchmark helper needs torchio I/O
+    preds, gt_paths = list(preds), list(gt_paths)
+    if len(preds) != len(gt_paths):
+        raise ValueError(f"preds and gt_paths must have equal length; got {len(preds)} and "
+                         f"{len(gt_paths)}.")
+
+    rows = []
+    for i, (pred, gt_path) in enumerate(zip(preds, gt_paths)):
+        gt = tio.LabelMap(gt_path)                       # single source: .data and .spacing agree
+        pred_data = tio.LabelMap(pred).data if isinstance(pred, (str, Path)) else pred
+        case_id = case_ids[i] if case_ids is not None else Path(str(gt_path)).name.split(".")[0]
+        row = {"case_id": case_id}
+        row.update(calculate_all_metrics(pred_data, gt.data, gt.spacing,
+                                         nsd_tolerances_mm=nsd_tolerances_mm, binary=binary))
+        rows.append(row)
+    return pd.DataFrame(rows)
+
 # %% ../nbs/05_vision_metrics.ipynb #js7clwse6n
 class AccumulatedDice(Metric):
     """nnU-Net-style accumulated Dice metric for reliable pseudo dice during training.