fastMONAI 0.5.2__py3-none-any.whl → 0.5.4__py3-none-any.whl

fastMONAI/vision_patch.py

@@ -0,0 +1,1125 @@
+"""Patch-based training and inference for 3D medical image segmentation using TorchIO's Queue mechanism."""
+
+# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/10_vision_patch.ipynb.
+
+# %% auto 0
+__all__ = ['normalize_patch_transforms', 'PatchConfig', 'med_to_subject', 'create_subjects_dataset', 'create_patch_sampler',
+           'MedPatchDataLoader', 'MedPatchDataLoaders', 'PatchInferenceEngine', 'patch_inference']
+
+# %% ../nbs/10_vision_patch.ipynb 2
+import torch
+import torchio as tio
+import pandas as pd
+import numpy as np
+import warnings
+from pathlib import Path
+from dataclasses import dataclass, field
+from typing import Callable
+from torch.utils.data import DataLoader
+from tqdm.auto import tqdm
+from fastai.data.all import *
+from .vision_core import MedImage, MedMask, MedBase, med_img_reader
+from .vision_inference import _to_original_orientation, _do_resize
+
+# %% ../nbs/10_vision_patch.ipynb 3
+def _get_default_device() -> torch.device:
+    """Get the default device (CUDA if available, else CPU)."""
+    return torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+def _warn_config_override(param_name: str, config_value, explicit_value):
+    """Warn when explicit argument overrides config value.
+
+    Args:
+        param_name: Name of the parameter (e.g., 'apply_reorder', 'target_spacing')
+        config_value: Value from PatchConfig
+        explicit_value: Explicitly provided value
+    """
+    if explicit_value is not None and config_value is not None:
+        if explicit_value != config_value:
+            warnings.warn(
+                f"{param_name} mismatch: explicit={explicit_value}, config={config_value}. "
+                f"Using explicit argument."
+            )
+
+
+def _extract_tio_transform(tfm):
+    """Extract TorchIO transform from fastMONAI wrapper or return as-is.
+
+    This function enables using fastMONAI wrappers (e.g., RandomAffine, RandomGamma)
+    in patch-based workflows where raw TorchIO transforms are needed for tio.Compose().
+
+    Uses the explicit `.tio_transform` property when available on fastMONAI wrappers.
+    Falls back to returning the transform unchanged for raw TorchIO transforms.
+
+    Args:
+        tfm: fastMONAI wrapper (e.g., RandomAffine) or raw TorchIO transform
+
+    Returns:
+        The underlying TorchIO transform
+
+    Example:
+        >>> from fastMONAI.vision_augmentation import RandomAffine
+        >>> wrapped = RandomAffine(degrees=10)
+        >>> raw = _extract_tio_transform(wrapped)  # Returns tio.RandomAffine
+    """
+    return getattr(tfm, 'tio_transform', tfm)
+
+
+def normalize_patch_transforms(tfms: list) -> list:
+    """Normalize transforms for patch-based workflow.
+
+    Extracts underlying TorchIO transforms from fastMONAI wrappers.
+    Also accepts raw TorchIO transforms for backward compatibility.
+
+    This enables using the same transform syntax in both standard and
+    patch-based workflows:
+
+        >>> from fastMONAI.vision_augmentation import RandomAffine, RandomGamma
+        >>>
+        >>> # Same syntax works in both contexts
+        >>> item_tfms = [RandomAffine(degrees=10), RandomGamma(p=0.5)]   # Standard
+        >>> patch_tfms = [RandomAffine(degrees=10), RandomGamma(p=0.5)]  # Patch-based
+
+    Args:
+        tfms: List of fastMONAI wrappers or raw TorchIO transforms
+
+    Returns:
+        List of raw TorchIO transforms suitable for tio.Compose()
+    """
+    if tfms is None:
+        return None
+    return [_extract_tio_transform(t) for t in tfms]
+
+# %% ../nbs/10_vision_patch.ipynb 7
+@dataclass
+class PatchConfig:
+    """Configuration for patch-based training and inference.
+    
+    Args:
+        patch_size: Size of patches [x, y, z].
+        patch_overlap: Overlap for inference GridSampler (int, float 0-1, or list).
+            - Float 0-1: fraction of patch_size (e.g., 0.5 = 50% overlap)
+            - Int >= 1: pixel overlap (e.g., 48 = 48 pixel overlap)
+            - List: per-dimension overlap in pixels
+        samples_per_volume: Number of patches to extract per volume during training.
+        sampler_type: Type of sampler ('uniform', 'label', 'weighted').
+        label_probabilities: For LabelSampler, dict mapping label values to probabilities.
+        queue_length: Maximum number of patches to store in queue.
+        queue_num_workers: Number of workers for parallel patch extraction.
+        aggregation_mode: For inference, how to combine overlapping patches ('crop', 'average', 'hann').
+        apply_reorder: Whether to reorder to RAS+ canonical orientation. Must match between
+            training and inference.
+        target_spacing: Target voxel spacing [x, y, z] for resampling. Must match between
+            training and inference.
+        padding_mode: Padding mode for CropOrPad when image < patch_size. Default is 0 (zero padding)
+            to align with nnU-Net's approach. Can be int, float, or string (e.g., 'minimum', 'mean').
+        keep_largest_component: If True, keep only the largest connected component
+            in binary segmentation predictions. Only applies during inference when
+            return_probabilities=False. Defaults to False.
+    
+    Example:
+        >>> config = PatchConfig(
+        ...     patch_size=[96, 96, 96],
+        ...     samples_per_volume=16,
+        ...     sampler_type='label',
+        ...     label_probabilities={0: 0.1, 1: 0.9},
+        ...     apply_reorder=True,
+        ...     target_spacing=[0.5, 0.5, 0.5]
+        ... )
+    """
+    patch_size: list = field(default_factory=lambda: [96, 96, 96])
+    patch_overlap: int | float | list = 0
+    samples_per_volume: int = 8
+    sampler_type: str = 'uniform'
+    label_probabilities: dict = None
+    queue_length: int = 300
+    queue_num_workers: int = 4
+    aggregation_mode: str = 'hann'
+    # Preprocessing parameters - must match between training and inference
+    apply_reorder: bool = False
+    target_spacing: list = None
+    padding_mode: int | float | str = 0  # Zero padding (nnU-Net standard)
+    # Post-processing (binary segmentation only)
+    keep_largest_component: bool = False
+    
+    def __post_init__(self):
+        """Validate configuration."""
+        valid_samplers = ['uniform', 'label', 'weighted']
+        if self.sampler_type not in valid_samplers:
+            raise ValueError(f"sampler_type must be one of {valid_samplers}")
+        
+        valid_aggregation = ['crop', 'average', 'hann']
+        if self.aggregation_mode not in valid_aggregation:
+            raise ValueError(f"aggregation_mode must be one of {valid_aggregation}")
+        
+        # Validate patch_overlap
+        # Negative overlap doesn't make sense
+        if isinstance(self.patch_overlap, (int, float)):
+            if self.patch_overlap < 0:
+                raise ValueError("patch_overlap cannot be negative")
+            # Check if overlap as pixels would exceed patch_size (causes step_size=0)
+            if self.patch_overlap >= 1:  # Pixel value, not fraction
+                for ps in self.patch_size:
+                    if self.patch_overlap >= ps:
+                        raise ValueError(
+                            f"patch_overlap ({self.patch_overlap}) must be less than patch_size ({ps}). "
+                            f"Overlap >= patch_size creates step_size <= 0 (infinite patches)."
+                        )
+        elif isinstance(self.patch_overlap, (list, tuple)):
+            for i, (overlap, ps) in enumerate(zip(self.patch_overlap, self.patch_size)):
+                if overlap < 0:
+                    raise ValueError(f"patch_overlap[{i}] cannot be negative")
+                if overlap >= ps:
+                    raise ValueError(
+                        f"patch_overlap[{i}] ({overlap}) must be less than patch_size[{i}] ({ps}). "
+                        f"Overlap >= patch_size creates step_size <= 0 (infinite patches)."
+                    )
+
+# %% ../nbs/10_vision_patch.ipynb 10
+def med_to_subject(
+    img: Path | str,
+    mask: Path | str = None,
+) -> tio.Subject:
+    """Create TorchIO Subject with LAZY loading (paths only, no tensor loading).
+    
+    This function stores file paths in the Subject, allowing TorchIO's Queue
+    workers to load volumes on-demand during training. This is memory-efficient
+    as volumes are not loaded into RAM until needed.
+    
+    Args:
+        img: Path to image file.
+        mask: Path to mask file (optional).
+    
+    Returns:
+        TorchIO Subject with 'image' and optionally 'mask' keys (lazy loaded).
+    
+    Example:
+        >>> subject = med_to_subject('image.nii.gz', 'mask.nii.gz')
+        >>> # Volume NOT loaded yet - only path stored
+        >>> data = subject['image'].data  # NOW volume is loaded
+    """
+    subject_dict = {
+        'image': tio.ScalarImage(path=str(img))  # Lazy - stores path only
+    }
+    
+    if mask is not None:
+        subject_dict['mask'] = tio.LabelMap(path=str(mask))  # Lazy
+    
+    return tio.Subject(**subject_dict)
+
+# %% ../nbs/10_vision_patch.ipynb 11
+def create_subjects_dataset(
+    df: pd.DataFrame,
+    img_col: str,
+    mask_col: str = None,
+    pre_tfms: list = None,
+    ensure_affine_consistency: bool = True
+) -> tio.SubjectsDataset:
+    """Build TorchIO SubjectsDataset with LAZY loading from DataFrame.
+
+    This function creates a SubjectsDataset that stores only file paths,
+    not loaded tensors. Volumes are loaded on-demand by Queue workers,
+    keeping memory usage constant regardless of dataset size.
+
+    Args:
+        df: DataFrame with image (and optionally mask) paths.
+        img_col: Column name containing image paths.
+        mask_col: Column name containing mask paths (optional).
+        pre_tfms: List of TorchIO transforms to apply before patch extraction.
+                  Use tio.ToCanonical() for reordering and tio.Resample() for resampling.
+        ensure_affine_consistency: If True and mask_col is provided, automatically
+            prepends tio.CopyAffine(target='image') to ensure spatial metadata
+            consistency between image and mask. This prevents "More than one value
+            for direction found" errors. Defaults to True.
+
+    Returns:
+        TorchIO SubjectsDataset with lazy-loaded subjects.
+
+    Example:
+        >>> # Preprocessing via transforms (applied by workers on-demand)
+        >>> pre_tfms = [
+        ...     tio.ToCanonical(),           # Reorder to RAS+
+        ...     tio.Resample([0.5, 0.5, 0.5]),  # Resample
+        ...     tio.ZNormalization(),        # Intensity normalization
+        ... ]
+        >>> dataset = create_subjects_dataset(
+        ...     df, img_col='image', mask_col='label',
+        ...     pre_tfms=pre_tfms
+        ... )
+        >>> # Memory: ~0 MB (only paths stored, not volumes)
+    """
+    subjects = []
+    for idx, row in df.iterrows():
+        img_path = row[img_col]
+        mask_path = row[mask_col] if mask_col else None
+
+        # Create subject with lazy loading (paths only)
+        subject = med_to_subject(img=img_path, mask=mask_path)
+        subjects.append(subject)
+
+    # Build transform pipeline
+    all_transforms = []
+
+    # Add CopyAffine as FIRST transform when mask is present
+    # This ensures spatial metadata consistency before other transforms
+    if mask_col is not None and ensure_affine_consistency:
+        all_transforms.append(tio.CopyAffine(target='image'))
+
+    # Add user-provided transforms
+    if pre_tfms:
+        all_transforms.extend(pre_tfms)
+
+    transform = tio.Compose(all_transforms) if all_transforms else None
+
+    return tio.SubjectsDataset(subjects, transform=transform)
+
+# %% ../nbs/10_vision_patch.ipynb 13
+def create_patch_sampler(config: PatchConfig) -> tio.data.PatchSampler:
+    """Create appropriate TorchIO sampler based on config.
+    
+    Args:
+        config: PatchConfig with sampler settings.
+    
+    Returns:
+        TorchIO PatchSampler instance.
+    
+    Example:
+        >>> config = PatchConfig(patch_size=[96, 96, 96], sampler_type='label')
+        >>> sampler = create_patch_sampler(config)
+    """
+    patch_size = config.patch_size
+    
+    if config.sampler_type == 'uniform':
+        return tio.UniformSampler(patch_size)
+    
+    elif config.sampler_type == 'label':
+        return tio.LabelSampler(
+            patch_size,
+            label_name='mask',
+            label_probabilities=config.label_probabilities
+        )
+    
+    elif config.sampler_type == 'weighted':
+        raise NotImplementedError(
+            "WeightedSampler requires a pre-computed probability map which is not currently supported. "
+            "Use 'label' sampler with label_probabilities for weighted sampling based on segmentation labels, "
+            "or 'uniform' for random patch extraction."
+        )
+    
+    raise ValueError(f"Unknown sampler type: {config.sampler_type}")
+
+# %% ../nbs/10_vision_patch.ipynb 17
+class MedPatchDataLoader:
+    """DataLoader wrapper for patch-based training with TorchIO Queue.
+
+    This class wraps a TorchIO Queue to provide a fastai-compatible DataLoader
+    interface for patch-based training.
+
+    Args:
+        subjects_dataset: TorchIO SubjectsDataset.
+        config: PatchConfig with queue and sampler settings.
+        batch_size: Number of patches per batch. Must be positive.
+        patch_tfms: Transforms to apply to extracted patches (training only).
+            Accepts both fastMONAI wrappers (e.g., RandomAffine, RandomGamma) and
+            raw TorchIO transforms. fastMONAI wrappers are automatically normalized
+            to raw TorchIO for internal use.
+        shuffle: Whether to shuffle subjects and patches.
+        drop_last: Whether to drop last incomplete batch.
+    """
+
+    def __init__(
+        self,
+        subjects_dataset: tio.SubjectsDataset,
+        config: PatchConfig,
+        batch_size: int = 4,
+        patch_tfms: list = None,
+        shuffle: bool = True,
+        drop_last: bool = False
+    ):
+        if batch_size <= 0:
+            raise ValueError(f"batch_size must be positive, got {batch_size}")
+        
+        self.subjects_dataset = subjects_dataset
+        self.config = config
+        self.bs = batch_size
+        self.shuffle = shuffle
+        self.drop_last = drop_last
+        self._device = _get_default_device()
+
+        # Create sampler
+        self.sampler = create_patch_sampler(config)
+
+        # Create patch transforms
+        # Normalize transforms - accepts both fastMONAI wrappers and raw TorchIO
+        normalized_tfms = normalize_patch_transforms(patch_tfms)
+        self.patch_tfms = tio.Compose(normalized_tfms) if normalized_tfms else None
+
+        # Create queue
+        self.queue = tio.Queue(
+            subjects_dataset,
+            max_length=config.queue_length,
+            samples_per_volume=config.samples_per_volume,
+            sampler=self.sampler,
+            num_workers=config.queue_num_workers,
+            shuffle_subjects=shuffle,
+            shuffle_patches=shuffle
+        )
+
+        # Create torch DataLoader
+        self._dl = DataLoader(
+            self.queue,
+            batch_size=batch_size,
+            num_workers=0,  # Queue handles workers
+            drop_last=drop_last
+        )
+
+    def __iter__(self):
+        """Iterate over batches, yielding (image, mask) tuples."""
+        for batch in self._dl:
+            # Extract image and mask tensors
+            img = batch['image'][tio.DATA]  # [B, C, H, W, D]
+            has_mask = 'mask' in batch
+
+            # Apply patch transforms if provided
+            if self.patch_tfms is not None:
+                # Apply transforms to each sample in batch
+                transformed_imgs = []
+                transformed_masks = [] if has_mask else None
+
+                for i in range(img.shape[0]):
+                    # Build subject dict with image, and mask if available
+                    subject_dict = {'image': tio.ScalarImage(tensor=batch['image'][tio.DATA][i])}
+                    if has_mask:
+                        subject_dict['mask'] = tio.LabelMap(tensor=batch['mask'][tio.DATA][i])
+
+                    subject = tio.Subject(subject_dict)
+                    transformed = self.patch_tfms(subject)
+                    transformed_imgs.append(transformed['image'].data)
+                    if has_mask:
+                        transformed_masks.append(transformed['mask'].data)
+
+                img = torch.stack(transformed_imgs)
+                mask = torch.stack(transformed_masks) if has_mask else None
+            else:
+                mask = batch['mask'][tio.DATA] if has_mask else None
+
+            # Convert to MedImage/MedMask and move to device
+            img = MedImage(img).to(self._device)
+            if mask is not None:
+                mask = MedMask(mask).to(self._device)
+
+            yield img, mask
+
+    def __len__(self):
+        """Return number of batches per epoch."""
+        n_patches = len(self.subjects_dataset) * self.config.samples_per_volume
+        if self.drop_last:
+            return n_patches // self.bs
+        return (n_patches + self.bs - 1) // self.bs
+
+    @property
+    def dataset(self):
+        """Return the underlying queue as dataset."""
+        return self.queue
+
+    @property
+    def device(self):
+        """Return current device."""
+        return self._device
+
+    def to(self, device):
+        """Move DataLoader to device."""
+        self._device = device
+        return self
+
+    def one_batch(self):
+        """Return one batch from the DataLoader.
+
+        Required for fastai compatibility - used for device detection
+        and batch shape validation during Learner initialization.
+
+        Returns:
+            Tuple of (image, mask) tensors on the correct device.
+        """
+        return next(iter(self))
+
+# %% ../nbs/10_vision_patch.ipynb 18
+class MedPatchDataLoaders:
+    """fastai-compatible DataLoaders for patch-based training with LAZY loading.
+
+    This class provides train and validation DataLoaders that work with
+    fastai's Learner for patch-based training on 3D medical images.
+
+    Memory-efficient: Volumes are loaded on-demand by Queue workers,
+    keeping memory usage constant (~150 MB) regardless of dataset size.
+
+    Note: Validation uses the same sampling as training (pseudo Dice).
+    For true validation metrics, use PatchInferenceEngine with GridSampler
+    for full-volume sliding window inference.
+
+    Example:
+        >>> import torchio as tio
+        >>>
+        >>> # New pattern: preprocessing params in config (DRY)
+        >>> config = PatchConfig(
+        ...     patch_size=[96, 96, 96],
+        ...     apply_reorder=True,
+        ...     target_spacing=[0.5, 0.5, 0.5]
+        ... )
+        >>> dls = MedPatchDataLoaders.from_df(
+        ...     df, img_col='image', mask_col='label',
+        ...     valid_pct=0.2,
+        ...     patch_config=config,
+        ...     pre_patch_tfms=[tio.ZNormalization()],
+        ...     bs=4
+        ... )
+        >>> learn = Learner(dls, model, loss_func=DiceLoss())
+    """
+
+    def __init__(
+        self,
+        train_dl: MedPatchDataLoader,
+        valid_dl: MedPatchDataLoader,
+        device: torch.device = None
+    ):
+        self._train_dl = train_dl
+        self._valid_dl = valid_dl
+        self._device = device or _get_default_device()
+
+        # Move to device
+        self._train_dl.to(self._device)
+        self._valid_dl.to(self._device)
+
+    @classmethod
+    def from_df(
+        cls,
+        df: pd.DataFrame,
+        img_col: str,
+        mask_col: str = None,
+        valid_pct: float = 0.2,
+        valid_col: str = None,
+        patch_config: PatchConfig = None,
+        pre_patch_tfms: list = None,
+        patch_tfms: list = None,
+        apply_reorder: bool = None,
+        target_spacing: list = None,
+        bs: int = 4,
+        seed: int = None,
+        device: torch.device = None,
+        ensure_affine_consistency: bool = True
+    ) -> 'MedPatchDataLoaders':
+        """Create train/valid DataLoaders from DataFrame with LAZY loading.
+
+        Memory-efficient: Only file paths are stored at creation time.
+        Volumes are loaded on-demand by Queue workers during training.
+
+        Note: Both train and valid use the same sampling strategy from patch_config.
+        This gives pseudo Dice during training. For true validation metrics,
+        use PatchInferenceEngine with full-volume sliding window inference.
+
+        Args:
+            df: DataFrame with image paths.
+            img_col: Column name for image paths.
+            mask_col: Column name for mask paths.
+            valid_pct: Fraction of data for validation.
+            valid_col: Column name for train/valid split (if pre-defined).
+            patch_config: PatchConfig instance. Preprocessing params (apply_reorder,
+                target_spacing) can be set here for DRY usage with PatchInferenceEngine.
+            pre_patch_tfms: TorchIO transforms applied before patch extraction
+                           (after reorder/resample). Example: [tio.ZNormalization()].
+            patch_tfms: TorchIO transforms applied to extracted patches (training only).
+            apply_reorder: If True, reorder to RAS+ orientation. If None, uses
+                patch_config.apply_reorder. Explicit value overrides config.
+            target_spacing: Target voxel spacing [x, y, z]. If None, uses
+                patch_config.target_spacing. Explicit value overrides config.
+            bs: Batch size.
+            seed: Random seed for splitting.
+            device: Device to use.
+            ensure_affine_consistency: If True and mask_col is provided, automatically
+                adds tio.CopyAffine(target='image') as the first transform to prevent
+                spatial metadata mismatch errors. Defaults to True.
+
+        Returns:
+            MedPatchDataLoaders instance.
+
+        Example:
+            >>> # New pattern: config contains preprocessing params
+            >>> config = PatchConfig(
+            ...     patch_size=[96, 96, 96],
+            ...     apply_reorder=True,
+            ...     target_spacing=[0.5, 0.5, 0.5],
+            ...     label_probabilities={0: 0.1, 1: 0.9}
+            ... )
+            >>> dls = MedPatchDataLoaders.from_df(
+            ...     df, img_col='image', mask_col='label',
+            ...     patch_config=config,
+            ...     pre_patch_tfms=[tio.ZNormalization()],
+            ...     patch_tfms=[tio.RandomAffine(degrees=10), tio.RandomFlip()],
+            ...     bs=4
+            ... )
+            >>> # Memory: ~150 MB (queue buffer only)
+        """
+        if patch_config is None:
+            patch_config = PatchConfig()
+
+        # Use config values, allow explicit overrides for backward compatibility
+        _apply_reorder = apply_reorder if apply_reorder is not None else patch_config.apply_reorder
+        _target_spacing = target_spacing if target_spacing is not None else patch_config.target_spacing
+
+        # Warn if both config and explicit args provided with different values
+        _warn_config_override('apply_reorder', patch_config.apply_reorder, apply_reorder)
+        _warn_config_override('target_spacing', patch_config.target_spacing, target_spacing)
+
+        # Split data
+        if valid_col is not None:
+            train_df = df[df[valid_col] == False].reset_index(drop=True)
+            valid_df = df[df[valid_col] == True].reset_index(drop=True)
+        else:
+            if seed is not None:
+                np.random.seed(seed)
+            n = len(df)
+            valid_idx = np.random.choice(n, size=int(n * valid_pct), replace=False)
+            train_idx = np.setdiff1d(np.arange(n), valid_idx)
+            train_df = df.iloc[train_idx].reset_index(drop=True)
+            valid_df = df.iloc[valid_idx].reset_index(drop=True)
+
+        # Build preprocessing transforms
+        all_pre_tfms = []
+
+        # Add reorder transform (reorder to RAS+ orientation)
+        if _apply_reorder:
+            all_pre_tfms.append(tio.ToCanonical())
+
+        # Add resample transform
+        if _target_spacing is not None:
+            all_pre_tfms.append(tio.Resample(_target_spacing))
+
+        # Add user-provided transforms
+        if pre_patch_tfms:
+            all_pre_tfms.extend(pre_patch_tfms)
+
+        # Create subjects datasets with lazy loading (paths only, ~0 MB)
+        train_subjects = create_subjects_dataset(
+            train_df, img_col, mask_col,
+            pre_tfms=all_pre_tfms if all_pre_tfms else None,
+            ensure_affine_consistency=ensure_affine_consistency
+        )
+        valid_subjects = create_subjects_dataset(
+            valid_df, img_col, mask_col,
+            pre_tfms=all_pre_tfms if all_pre_tfms else None,
+            ensure_affine_consistency=ensure_affine_consistency
+        )
+
+        # Create DataLoaders (both use same patch_config for consistent sampling)
+        train_dl = MedPatchDataLoader(
+            train_subjects, patch_config, bs,
+            patch_tfms=patch_tfms, shuffle=True, drop_last=True
+        )
+        valid_dl = MedPatchDataLoader(
+            valid_subjects, patch_config, bs,
+            patch_tfms=None,  # No augmentation for validation
+            shuffle=False, drop_last=False
+        )
+
+        # Create instance and store metadata
+        instance = cls(train_dl, valid_dl, device)
+        instance._img_col = img_col
+        instance._mask_col = mask_col
+        instance._pre_patch_tfms = pre_patch_tfms
+        instance._apply_reorder = _apply_reorder
+        instance._target_spacing = _target_spacing
+        instance._ensure_affine_consistency = ensure_affine_consistency
+        instance._patch_config = patch_config
+        return instance
+
+    @property
+    def train(self):
+        """Training DataLoader."""
+        return self._train_dl
+
+    @property
+    def valid(self):
+        """Validation DataLoader."""
+        return self._valid_dl
+
+    @property
+    def train_ds(self):
+        """Training subjects dataset."""
+        return self._train_dl.subjects_dataset
+
+    @property
+    def valid_ds(self):
+        """Validation subjects dataset."""
+        return self._valid_dl.subjects_dataset
+
+    @property
+    def device(self):
+        """Current device."""
+        return self._device
+
+    @property
+    def bs(self):
+        """Batch size."""
+        return self._train_dl.bs
+
+    @property
+    def apply_reorder(self):
+        """Whether reordering to RAS+ is enabled."""
+        return getattr(self, '_apply_reorder', False)
+
+    @property
+    def target_spacing(self):
+        """Target voxel spacing for resampling."""
+        return getattr(self, '_target_spacing', None)
+
+    @property
+    def patch_config(self):
+        """The PatchConfig used for this DataLoaders."""
+        return getattr(self, '_patch_config', None)
+
+    def to(self, device):
+        """Move DataLoaders to device."""
+        self._device = device
+        self._train_dl.to(device)
+        self._valid_dl.to(device)
+        return self
+
+    def __iter__(self):
+        """Iterate over training DataLoader."""
+        return iter(self._train_dl)
+
+    def one_batch(self):
+        """Return one batch from the training DataLoader.
+
+        Required for fastai Learner compatibility - used for device
+        detection and batch shape validation.
+        """
+        return self._train_dl.one_batch()
+
+    def __len__(self):
+        """Return number of batches in training DataLoader."""
+        return len(self._train_dl)
+
+    def __getitem__(self, idx):
+        """Get DataLoader by index. Required for fastai Learner compatibility.
+
+        Args:
+            idx: 0 for training DataLoader, 1 for validation DataLoader.
+
+        Returns:
+            MedPatchDataLoader instance.
+        """
+        if idx == 0:
+            return self._train_dl
+        elif idx == 1:
+            return self._valid_dl
+        else:
+            raise IndexError(f"Index {idx} out of range. Use 0 (train) or 1 (valid).")
+
+    def cuda(self):
+        """Move DataLoaders to CUDA device."""
+        return self.to(torch.device('cuda'))
+
+    def cpu(self):
+        """Move DataLoaders to CPU."""
+        return self.to(torch.device('cpu'))
+
+    def new_empty(self):
+        """Create a new empty version of self for learner export.
+
+        Required for fastai Learner.export() compatibility - creates a
+        lightweight placeholder that can be pickled without the full dataset.
+
+        Returns:
+            A minimal MedPatchDataLoaders-like object with no data.
+        """
+        class EmptyMedPatchDataLoaders:
+            """Minimal placeholder for exported learner."""
+            def __init__(self, device):
+                self._device = device
+            @property
+            def device(self): return self._device
+            def to(self, device):
+                self._device = device
+                return self
+
+        return EmptyMedPatchDataLoaders(self._device)
+
+# %% ../nbs/10_vision_patch.ipynb 20
+import numbers
+
+def _normalize_patch_overlap(patch_overlap, patch_size):
+    """Convert patch_overlap to integer pixel values for TorchIO compatibility.
+
+    TorchIO's GridSampler expects patch_overlap as a tuple of even integers.
+    This function handles:
+    - Fractional overlap (0-1): converted to pixel values based on patch_size
+    - Numpy scalar types: converted to native Python types
+    - Sequences: converted to tuple of integers
+
+    Note: Input validation (negative values, overlap >= patch_size) is handled
+    by PatchConfig.__post_init__(). This function focuses on format conversion.
+
+    Args:
+        patch_overlap: int, float (0-1 for fraction), or sequence
+        patch_size: list/tuple of patch dimensions [x, y, z]
+
+    Returns:
+        Tuple of even integers suitable for TorchIO GridSampler
+    """
+    # Handle scalar fractional overlap (0 < x < 1)
+    # Note: excludes 1.0 as 100% overlap creates step_size=0 (infinite patches)
+    if isinstance(patch_overlap, (int, float, numbers.Number)) and 0 < float(patch_overlap) < 1:
+        # Convert fraction to pixel values, ensure even
+        result = []
+        for ps in patch_size:
+            pixels = int(int(ps) * float(patch_overlap))
+            # Ensure even (required by TorchIO)
+            if pixels % 2 != 0:
+                pixels = pixels - 1 if pixels > 0 else 0
+            result.append(pixels)
+        return tuple(result)
+
+    # Handle scalar integer (including numpy scalars) - values > 1 are pixel counts
+    if isinstance(patch_overlap, (int, float, numbers.Number)):
+        val = int(patch_overlap)
+        # Ensure even
+        if val % 2 != 0:
+            val = val - 1 if val > 0 else 0
+        return tuple(val for _ in patch_size)
+
+    # Handle sequences (list, tuple, ndarray)
+    result = []
+    for val in patch_overlap:
+        pixels = int(val)
+        if pixels % 2 != 0:
+            pixels = pixels - 1 if pixels > 0 else 0
+        result.append(pixels)
+    return tuple(result)
+
+
+class PatchInferenceEngine:
+    """Patch-based inference with automatic volume reconstruction.
+    
+    Uses TorchIO's GridSampler to extract overlapping patches and
+    GridAggregator to reconstruct the full volume from predictions.
+    
+    Args:
+        learner: PyTorch model or fastai Learner.
+        config: PatchConfig with inference settings. Preprocessing params (apply_reorder,
+            target_spacing, padding_mode) can be set here for DRY usage.
+        apply_reorder: Whether to reorder to RAS+ orientation. If None, uses config value.
+        target_spacing: Target voxel spacing. If None, uses config value.
+        batch_size: Number of patches to predict at once. Must be positive.
+        pre_inference_tfms: List of TorchIO transforms to apply before patch extraction.
+            IMPORTANT: Should match the pre_patch_tfms used during training (e.g., [tio.ZNormalization()]).
+            This ensures preprocessing consistency between training and inference.
+    
+    Example:
+        >>> # DRY pattern: use same config for training and inference
+        >>> config = PatchConfig(
+        ...     patch_size=[96, 96, 96],
+        ...     apply_reorder=True,
+        ...     target_spacing=[1.0, 1.0, 1.0]
+        ... )
+        >>> # Training
+        >>> dls = MedPatchDataLoaders.from_df(df, 'img', 'mask', patch_config=config)
+        >>> # Inference - no need to repeat params!
+        >>> engine = PatchInferenceEngine(
+        ...     learn, config,
+        ...     pre_inference_tfms=[tio.ZNormalization()]
+        ... )
+        >>> pred = engine.predict('image.nii.gz')
+    """
+    
+    def __init__(
+        self,
+        learner,
+        config: PatchConfig,
+        apply_reorder: bool = None,
+        target_spacing: list = None,
+        batch_size: int = 4,
+        pre_inference_tfms: list = None
+    ):
+        if batch_size <= 0:
+            raise ValueError(f"batch_size must be positive, got {batch_size}")
+        
+        # Extract model from Learner if needed
+        self.model = learner.model if hasattr(learner, 'model') else learner
+        self.config = config
+        self.batch_size = batch_size
+        self.pre_inference_tfms = tio.Compose(pre_inference_tfms) if pre_inference_tfms else None
+        
+        # Use config values, allow explicit overrides for backward compatibility
+        self.apply_reorder = apply_reorder if apply_reorder is not None else config.apply_reorder
+        self.target_spacing = target_spacing if target_spacing is not None else config.target_spacing
+        
+        # Warn if explicit args provided but differ from config (potential mistake)
+        _warn_config_override('apply_reorder', config.apply_reorder, apply_reorder)
+        _warn_config_override('target_spacing', config.target_spacing, target_spacing)
+        
+        # Get device from model parameters, with fallback for parameter-less models
+        try:
+            self._device = next(self.model.parameters()).device
+        except StopIteration:
+            self._device = _get_default_device()
+    
+    def predict(
+        self,
+        img_path: Path | str,
+        return_probabilities: bool = False,
+        return_affine: bool = False
+    ) -> torch.Tensor | tuple[torch.Tensor, np.ndarray]:
+        """Predict on a single volume using patch-based inference.
+
+        Args:
+            img_path: Path to input image.
+            return_probabilities: If True, return probability map instead of argmax.
+            return_affine: If True, return (prediction, affine) tuple instead of just prediction.
+
+        Returns:
+            Predicted segmentation mask tensor, or tuple (prediction, affine) if return_affine=True.
+        """
+        # Load image - keep org_img and org_size for post-processing
+        # Note: med_img_reader handles reorder/resample internally, no global state needed
+        org_img, input_img, org_size = med_img_reader(
+            img_path, apply_reorder=self.apply_reorder, target_spacing=self.target_spacing, only_tensor=False
+        )
+
+        # Create TorchIO Subject from preprocessed image
+        subject = tio.Subject(
+            image=tio.ScalarImage(tensor=input_img.data.float(), affine=input_img.affine)
+        )
+
+        # Apply pre-inference transforms (e.g., ZNormalization) to match training
+        if self.pre_inference_tfms is not None:
+            subject = self.pre_inference_tfms(subject)
+
+        # Pad dimensions smaller than patch_size, keep larger dimensions intact
+        # GridSampler handles large images via overlapping patches
+        img_shape = subject['image'].shape[1:]  # Exclude channel dim
+        target_size = [max(s, p) for s, p in zip(img_shape, self.config.patch_size)]
+        
+        # Warn if volume needed padding (may cause artifacts if training didn't cover similar sizes)
+        if any(s < p for s, p in zip(img_shape, self.config.patch_size)):
+            padded_dims = [f"dim{i}: {s}<{p}" for i, (s, p) in enumerate(zip(img_shape, self.config.patch_size)) if s < p]
+            warnings.warn(
+                f"Image size {list(img_shape)} smaller than patch_size {self.config.patch_size} "
+                f"in {padded_dims}. Padding with mode={self.config.padding_mode}. "
+                "Ensure training data covered similar sizes to avoid artifacts."
+            )
+        
+        # Use padding_mode from config (default: 0 for zero padding, nnU-Net standard)
+        subject = tio.CropOrPad(target_size, padding_mode=self.config.padding_mode)(subject)
+
+        # Convert patch_overlap to integer pixel values for TorchIO compatibility
+        patch_overlap = _normalize_patch_overlap(self.config.patch_overlap, self.config.patch_size)
+
+        # Create GridSampler
+        grid_sampler = tio.GridSampler(
+            subject,
+            patch_size=self.config.patch_size,
+            patch_overlap=patch_overlap
+        )
+
+        # Create GridAggregator
+        aggregator = tio.GridAggregator(
+            grid_sampler,
+            overlap_mode=self.config.aggregation_mode
+        )
+
+        # Create patch loader
+        patch_loader = DataLoader(
+            grid_sampler,
+            batch_size=self.batch_size,
+            num_workers=0
+        )
+
+        # Predict patches
+        self.model.eval()
+        with torch.no_grad():
+            for patches_batch in patch_loader:
+                patch_input = patches_batch['image'][tio.DATA].to(self._device)
+                locations = patches_batch[tio.LOCATION]
+
+                # Forward pass - get logits
+                logits = self.model(patch_input)
+
+                # Convert logits to probabilities BEFORE aggregation
+                # This is critical: softmax is non-linear, so we must aggregate
+                # probabilities, not logits, to get correct boundary handling
+                n_classes = logits.shape[1]
+                if n_classes == 1:
+                    probs = torch.sigmoid(logits)
+                else:
+                    probs = torch.softmax(logits, dim=1)  # dim=1 for batch [B, C, H, W, D]
+
+                # Add probabilities to aggregator
+                aggregator.add_batch(probs.cpu(), locations)
+
+        # Get reconstructed output (now contains probabilities, not logits)
+        output = aggregator.get_output_tensor()
+
+        # Convert to prediction mask (only if not returning probabilities)
+        if return_probabilities:
+            result = output  # Keep as float probabilities
+        else:
+            n_classes = output.shape[0]
+            if n_classes == 1:
+                result = (output > 0.5).float()
+            else:
+                result = output.argmax(dim=0, keepdim=True).float()
+
+        # Apply keep_largest post-processing for binary segmentation
+        if not return_probabilities and self.config.keep_largest_component:
+            from fastMONAI.vision_inference import keep_largest
+            result = keep_largest(result.squeeze(0)).unsqueeze(0)
+
+        # Post-processing: resize back to original size and reorient
+        # This matches the workflow in vision_inference.py
+        
+        # Wrap result in TorchIO Image for resizing
+        # Use ScalarImage for probabilities, LabelMap for masks
+        if return_probabilities:
+            pred_img = tio.ScalarImage(tensor=result.float(), affine=input_img.affine)
+        else:
+            pred_img = tio.LabelMap(tensor=result.float(), affine=input_img.affine)
+        
+        # Resize back to original size (before resampling)
+        pred_img = _do_resize(pred_img, org_size, image_interpolation='nearest')
+        
+        # Reorient to original orientation (if reorder was applied)
+        # Use explicit .cpu() for consistent device handling
+        if self.apply_reorder:
+            reoriented_array = _to_original_orientation(
+                pred_img.as_sitk(),
+                ('').join(org_img.orientation)
+            )
+            result = torch.from_numpy(reoriented_array).cpu()
+            # Only convert to long for masks, not probabilities
+            if not return_probabilities:
+                result = result.long()
+        else:
+            result = pred_img.data.cpu()
+            # Only convert to long for masks, not probabilities
+            if not return_probabilities:
+                result = result.long()
+
+        # Use original affine matrix for correct spatial alignment
+        # org_img.affine is always available from med_img_reader
+        if not (hasattr(org_img, 'affine') and org_img.affine is not None):
+            raise RuntimeError(
+                "org_img.affine not available. This should never happen - please report this bug."
+            )
+        affine = org_img.affine.copy()
+
+        if return_affine:
+            return result, affine
+        return result
+    
+    def to(self, device):
+        """Move engine to device."""
+        self._device = device
+        self.model.to(device)
+        return self
+
+# %% ../nbs/10_vision_patch.ipynb 21
+def patch_inference(
+    learner,
+    config: PatchConfig,
+    file_paths: list,
+    apply_reorder: bool = None,
+    target_spacing: list = None,
+    batch_size: int = 4,
+    return_probabilities: bool = False,
+    progress: bool = True,
+    save_dir: str = None,
+    pre_inference_tfms: list = None
+) -> list:
+    """Batch patch-based inference on multiple volumes.
+    
+    Args:
+        learner: PyTorch model or fastai Learner.
+        config: PatchConfig with inference settings. Preprocessing params (apply_reorder,
+            target_spacing) can be set here for DRY usage.
+        file_paths: List of image paths.
+        apply_reorder: Whether to reorder to RAS+ orientation. If None, uses config value.
+        target_spacing: Target voxel spacing. If None, uses config value.
+        batch_size: Patches per batch.
+        return_probabilities: Return probability maps.
+        progress: Show progress bar.
+        save_dir: Directory to save predictions as NIfTI files. If None, predictions are not saved.
+        pre_inference_tfms: List of TorchIO transforms to apply before patch extraction.
+            IMPORTANT: Should match the pre_patch_tfms used during training (e.g., [tio.ZNormalization()]).
+    
+    Returns:
+        List of predicted tensors.
+    
+    Example:
+        >>> # DRY pattern: use same config for training and inference
+        >>> config = PatchConfig(
+        ...     patch_size=[96, 96, 96],
+        ...     apply_reorder=True,
+        ...     target_spacing=[0.4102, 0.4102, 1.5]
+        ... )
+        >>> predictions = patch_inference(
+        ...     learner=learn,
+        ...     config=config,  # apply_reorder and target_spacing from config
+        ...     file_paths=val_paths,
+        ...     pre_inference_tfms=[tio.ZNormalization()],
+        ...     save_dir='predictions/patch_based'
+        ... )
+    """
+    # Use config values if not explicitly provided
+    _apply_reorder = apply_reorder if apply_reorder is not None else config.apply_reorder
+    _target_spacing = target_spacing if target_spacing is not None else config.target_spacing
+    
+    engine = PatchInferenceEngine(
+        learner, config, _apply_reorder, _target_spacing, batch_size, pre_inference_tfms
+    )
+    
+    # Create save directory if specified
+    if save_dir is not None:
+        save_path = Path(save_dir)
+        save_path.mkdir(parents=True, exist_ok=True)
+    
+    predictions = []
+    iterator = tqdm(file_paths, desc='Patch inference') if progress else file_paths
+    
+    for path in iterator:
+        # Get prediction and affine when saving is needed
+        if save_dir is not None:
+            pred, affine = engine.predict(path, return_probabilities, return_affine=True)
+        else:
+            pred = engine.predict(path, return_probabilities)
+        predictions.append(pred)
+        
+        # Save prediction if save_dir specified
+        if save_dir is not None:
+            input_path = Path(path)
+            # Create output filename based on input using suffix-based approach
+            # This handles .nii.gz correctly without corrupting filenames with .nii elsewhere
+            stem = input_path.stem
+            if input_path.suffix == '.gz' and stem.endswith('.nii'):
+                # Handle .nii.gz files: stem is "filename.nii", strip the .nii
+                stem = stem[:-4]
+                out_name = f"{stem}_pred.nii.gz"
+            elif input_path.suffix == '.nii':
+                # Handle .nii files
+                out_name = f"{stem}_pred.nii"
+            else:
+                # Fallback for other formats
+                out_name = f"{stem}_pred.nii.gz"
+            out_path = save_path / out_name
+            
+            # affine is guaranteed to be valid from engine.predict() with return_affine=True
+            # Save as NIfTI using TorchIO with correct type
+            # Use ScalarImage for probabilities (float), LabelMap for masks (int)
+            if return_probabilities:
+                pred_img = tio.ScalarImage(tensor=pred, affine=affine)
+            else:
+                pred_img = tio.LabelMap(tensor=pred, affine=affine)
+            pred_img.save(out_path)
+    
+    return predictions