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

This tutorial demonstrates when and how to use patch-based training:

• Large volumes: When full images don’t fit in GPU memory
• Memory constraints: Constant ~150 MB memory usage regardless of dataset size
• Class imbalance: Foreground-weighted sampling for better learning on small structures

We use the same Heart MRI dataset from the binary segmentation tutorial for direct comparison.

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 takes the same parameters as patch_tfms, and each value means the same thing (per-axis values follow array axis order, lengths are in voxels), but the output is not bit-identical. Rotations act in normalised patch coordinates, so on a non-cubic patch a rotation between two axes of different length is stretched by their length ratio. Blur uses one isotropic sigma per sample, anisotropy resamples with F.interpolate, and gamma draws one exponent per sample instead of one per channel. 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()