Patch-Based Binary Segmentation

Patch-based training for 3D medical image segmentation using fastMONAI’s MedPatchDataLoaders with lazy loading, where volumes are loaded on-demand, keeping memory usage constant regardless of dataset size.

Google Colab

from fastMONAI.vision_all import *

from monai.apps import DecathlonDataset
from sklearn.model_selection import train_test_split
import torchio as tio

Download external data

We use the MONAI function DecathlonDataset to download the Heart MRI dataset from the Medical Segmentation Decathlon challenge.

path = Path('../data')
path.mkdir(exist_ok=True)

task = "Task02_Heart"
training_data = DecathlonDataset(root_dir=path, task=task, section="training", 
    download=True, cache_num=0, num_workers=3)

df = pd.DataFrame(training_data.data)
df.shape

Split the labeled data into training and test sets.

train_df, test_df = train_test_split(df, test_size=0.1, random_state=42)
train_df.shape, test_df.shape

Analyze training data

Use MedDataset to analyze the dataset and get preprocessing recommendations.

med_dataset = MedDataset(img_list=train_df.label.tolist(), dtype=MedMask, max_workers=12)

med_dataset.df.head()

data_info_df = med_dataset.summary()

suggestion = med_dataset.get_suggestion()
target_spacing = suggestion['target_spacing']
target_spacing

stats = med_dataset.get_size_statistics(target_spacing=target_spacing)
print(f"Image sizes (after resampling to {target_spacing}):")
print(f"  Min:    {stats['min']}")
print(f"  Median: {stats['median']}")
print(f"  Max:    {stats['max']}")

suggested_size = suggest_patch_size(med_dataset, target_spacing=target_spacing)
print(f"\nSuggested patch size: {suggested_size}")

Configure patch-based training

PatchConfig centralizes all patch-related parameters:

• patch_size: Size of extracted patches [x, y, z], should be divisible by 16 for UNet compatibility
• samples_per_volume: Number of patches extracted per volume per epoch
• sampler_type: 'uniform' (random) or 'label' (foreground-weighted)
• label_probabilities: For 'label' sampler, probability of sampling each class
• queue_length: Number of patches to keep in memory buffer
• patch_overlap: Overlap for inference (float 0-1 for fraction, or int for pixels)
• aggregation_mode: How to combine overlapping patches ('hann' for smooth boundaries)
• padding_mode: Padding mode when image < patch_size (0 = zero padding, nnU-Net standard)

patch_config = PatchConfig(
    patch_size=[160, 160, 80],
    samples_per_volume=8,
    sampler_type='label',
    label_probabilities={0: 0.5, 1: 0.5},
    patch_overlap=0.5,
    keep_largest_component=True,
    target_spacing=target_spacing,
    aggregation_mode='hann',
    # normalization is set once on the config and applied at both training and inference
    # (no need to re-specify it at inference time).
    normalization=[ZNormalization(masking_method='foreground')]
)

print(f"Patch config: {patch_config}")

Alternatively, use PatchConfig.from_dataset(med_dataset) to auto-configure the patch size from the dataset analysis.

Define transforms

Normalization is set on the PatchConfig above (normalization=...) and applied at both training and inference, so training and inference use the same preprocessing.

The remaining stage is patch_tfms: augmentations applied to extracted patches during training only.

suggest_patch_size returns a size that fits every volume, so no padding is needed at training time. Here we pick [160, 160, 80] manually; it is also smaller than every volume, so it needs no padding either.

You can still pass pre_patch_tfms= to from_df (or pre_inference_tfms= at inference) to override the config, which is useful for transforms that are not serializable.

patch_tfms = [
    RandomAffine(scales=(0.7, 1.4), degrees=30, translation=(25, 25, 10), p=0.2),
    RandomAnisotropy(downsampling=(1.5, 3), p=0.25),
    RandomGamma(log_gamma=(-0.3, 0.3), p=0.3),
    RandomIntensityScale(scale_range=(0.75, 1.25), p=0.1),
    RandomNoise(std=0.1, p=0.1),
    RandomBlur(std=(0.5, 1.0), p=0.2),
    RandomFlip(axes='LRAPIS', p=0.5),
]

Alternative: GPU-batched augmentation

For long training runs (e.g., hundreds of epochs), GPU-batched augmentation can reduce training time by moving transforms from CPU to GPU. Instead of per-sample TorchIO transforms (patch_tfms), gpu_patch_augmentations creates a batched augmentation pipeline that operates on GPU tensors directly.

Because it is genuinely batched (one grid_sample over the whole batch), it is far faster than per-sample augmentation. Measured on an RTX 6000 Ada (full pipeline, 128^3): ~3 ms/batch vs ~77 ms for the equivalent MONAI per-sample-on-GPU pipeline (~25x) and ~287 ms for the TorchIO CPU path (~100x).

It is an approximate reimplementation of the patch_tfms transforms (voxel-space affine, per-axis flip, image-only anisotropy, sign-preserving gamma), It approximates the CPU path rather than matching it exactly. Use gpu_augmentation instead of patch_tfms (they are mutually exclusive).

# # GPU augmentation alternative (uncomment to use instead of patch_tfms above)
# gpu_aug = gpu_patch_augmentations(patch_config.patch_size, patch_config.target_spacing)
#
# # Then pass gpu_augmentation instead of patch_tfms to from_df:
# # dls = MedPatchDataLoaders.from_df(
# #     ..., gpu_augmentation=gpu_aug, ...  # replaces patch_tfms=patch_tfms
# # )

Create patch-based DataLoaders

MedPatchDataLoaders uses lazy loading: - Only file paths are stored at creation time (~0 MB) - Volumes are loaded on-demand by Queue workers

bs = 4

dls = MedPatchDataLoaders.from_df(
    df=train_df,
    img_col='image',
    mask_col='label',
    valid_pct=0.1,
    patch_config=patch_config,
    patch_tfms=patch_tfms,
    bs=bs,
    seed=42
)

print(f"Training subjects: {len(dls.train.subjects_dataset)}")
print(f"Validation subjects: {len(dls.valid.subjects_dataset)}")

batch = next(iter(dls.train))
x, y = batch
print(f"Batch shape - Image: {x.shape}, Mask: {y.shape}")

dls.show_batch(anatomical_plane=2, max_n=2, overlay=False)

Create and train a 3D model

We use MONAI’s UNet with: - out_channels=2: Softmax output (background + foreground) - Instance normalization: Common in medical imaging - DiceCELoss: Combines Dice loss with Cross-Entropy for stable training

from monai.networks.nets import UNet
from monai.networks.layers import Norm
from monai.losses import DiceCELoss

• DiceCELoss: Combines Dice loss with Cross-Entropy for stable training. batch=True computes Dice by pooling TP/FP/FN across all samples in the batch before computing the ratio (nnU-Net default), rather than averaging per-sample Dice scores. This provides more stable gradients when some patches contain little or no foreground.

model = UNet(
    spatial_dims=3,
    in_channels=1,
    out_channels=2,
    channels=(16, 32, 64, 128, 256),
    strides=(2, 2, 2, 2),
    num_res_units=2,
    norm=Norm.INSTANCE
)

loss_func = CustomLoss(loss_func=DiceCELoss(
    to_onehot_y=True,
    softmax=True,
    include_background=False,
    batch=True
))

We use AccumulatedDice metric which accumulates true positives, false positives, and false negatives across all validation batches before computing Dice.

learn = Learner(dls, model, loss_func=loss_func, metrics=[AccumulatedDice(n_classes=2)])

learn.lr_find()

lr = 1e-3

best_model_fname = "best_heart_patch"
save_best = EMACheckpoint(
    monitor='accumulated_dice',
    momentum=0.9,
    comp=np.greater,
    fname=best_model_fname,
    with_opt=False
)

model_spec = make_model_spec('monai.unet', {
    'spatial_dims': 3, 'in_channels': 1, 'out_channels': 2,
    'channels': [16, 32, 64, 128, 256], 'strides': [2, 2, 2, 2],
    'num_res_units': 2, 'norm': 'INSTANCE',
})
output_spec = make_output_spec('multiclass_segmentation', classes=2)

mlflow_callback = create_mlflow_callback(
    learn, experiment_name="Task02_Heart_Patch",
    dataset_version=med_dataset.fingerprint,
    model_spec=model_spec, output_spec=output_spec)
learn.fit_one_cycle(40, lr, cbs=[mlflow_callback, save_best])

learn.recorder.plot_loss();

Evaluate on validation set

Evaluate the trained model on the validation set using sliding-window patch inference.

Metrics: - DSC: Dice score, overlap similarity (higher = better) - HD95: 95th percentile Hausdorff distance in mm (lower = better) - Sens: Sensitivity, the true-positive rate (higher = better) - LDR: Lesion detection rate (higher = better) - RVE: Relative volume error (0 = optimal, + = over-seg, - = under-seg)

val_subjects = dls.valid.subjects_dataset
val_img_paths = [str(s['image'].path) for s in val_subjects]
val_mask_paths = [str(s['mask'].path) for s in val_subjects]

print(f"Validation set: {len(val_img_paths)} images")

learn.load(best_model_fname)
print(f"Loaded best model: {best_model_fname}")

Path('models').mkdir(exist_ok=True)
save_safetensors_model(
    learn.model, 'models/best_model.safetensors', model_spec,
    patch_inference_config(patch_config, output_spec), artifact_role='best')

save_dir = Path('predictions/patch_heart_val')
predictions = patch_inference(
    learner=learn, config=patch_config, file_paths=val_img_paths,
    save_dir=str(save_dir), progress=True, tta=True)
print(f"\nSaved {len(predictions)} predictions to {save_dir}/")

from fastMONAI.vision_metrics import (calculate_dsc, calculate_surface_metrics,
                                       calculate_confusion_metrics,
                                       calculate_lesion_detection_rate,
                                       calculate_signed_rve)

results = []

for i, pred in enumerate(predictions):
    img_name = Path(val_img_paths[i]).name
    gt_path = val_mask_paths[i]

    # Ground truth and prediction are both in native voxel space (patch_inference resizes the
    # prediction back to the original grid), so load the GT natively too.
    gt = MedMask.create(gt_path)

    pred_5d = pred.unsqueeze(0).float()
    gt_5d = gt.data.unsqueeze(0).float()

    dsc = calculate_dsc(pred_5d, gt_5d).mean().item()
    # Spacing-aware HD95 in mm: use the per-case voxel spacing read from the GT file.
    # Predictions and GT live in native space, so the annotation's own resolution is correct.
    hd95 = calculate_surface_metrics(pred_5d, gt_5d, spacing_mm=tio.LabelMap(gt_path).spacing)['hd95_mm']
    sens = calculate_confusion_metrics(pred_5d, gt_5d, "sensitivity").nanmean().item()
    ldr = calculate_lesion_detection_rate(pred_5d, gt_5d).nanmean().item()
    rve = calculate_signed_rve(pred_5d, gt_5d).nanmean().item()

    results.append({
        'image': img_name, 'dsc': dsc, 'hd95': hd95,
        'sensitivity': sens, 'ldr': ldr, 'rve': rve
    })

results_df = pd.DataFrame(results)

mlflow_callback.log_metrics_table(results_df, display=True) 
mlflow_callback.log_metrics(
    {f'val_{m}': results_df[m].mean() for m in results_df.select_dtypes(include='number').columns}
)
mlflow_callback.log_dataframe(results_df)

from fastMONAI.vision_plot import show_segmentation_comparison

idx = 0
val_img = MedImage.create(
    val_img_paths[idx],
    apply_reorder=patch_config.apply_reorder,
    target_spacing=patch_config.target_spacing
)
val_gt = MedMask.create(
    val_mask_paths[idx],
    apply_reorder=patch_config.apply_reorder,
    target_spacing=patch_config.target_spacing
)

show_segmentation_comparison(
    image=val_img,
    ground_truth=val_gt,
    prediction=predictions[idx],
    metric_value=results_df.iloc[idx]['dsc'],
    voxel_size=patch_config.target_spacing,
    anatomical_plane=2  # axial view
)

View experiment tracking

mlflow_ui = MLflowUIManager()
mlflow_ui.start_ui()

Summary

In this tutorial, we demonstrated patch-based training and evaluation for 3D medical image segmentation:

Training: 1. PatchConfig: Centralized configuration for patch size, sampling, and inference parameters 2. MedPatchDataLoaders: Memory-efficient lazy loading with TorchIO Queue 3. Transforms: Full-volume normalization plus patch-only training augmentation 4. AccumulatedDice: nnU-Net-style accumulated validation metric 5. Safe artifacts: MLflow retains .pth checkpoints and adds strict-loadable .safetensors; no new Learner pickle is written

Evaluation: 6. patch_inference(): Batch sliding-window inference with NIfTI output 7. Complete config: preprocessing, padding_mode, binary_threshold, and output decoding are serialized explicitly

mlflow_ui.stop()