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

fastMONAI/vision_augmentation.py

@@ -1,16 +1,17 @@
 # AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/03_vision_augment.ipynb.
 
 # %% auto 0
-__all__ = ['CustomDictTransform', 'do_pad_or_crop', 'PadOrCrop', 'ZNormalization', 'RescaleIntensity', 'NormalizeIntensity',
-           'BraTSMaskConverter', 'BinaryConverter', 'RandomGhosting', 'RandomSpike', 'RandomNoise', 'RandomBiasField',
-           'RandomBlur', 'RandomGamma', 'RandomIntensityScale', 'RandomMotion', 'RandomElasticDeformation',
-           'RandomAffine', 'RandomFlip', 'OneOf']
+__all__ = ['CustomDictTransform', 'do_pad_or_crop', 'PadOrCrop', 'SpatialPad', 'ZNormalization', 'RescaleIntensity',
+           'NormalizeIntensity', 'BraTSMaskConverter', 'BinaryConverter', 'RandomGhosting', 'RandomSpike',
+           'RandomNoise', 'RandomBiasField', 'RandomBlur', 'RandomGamma', 'RandomIntensityScale', 'RandomMotion',
+           'RandomCutout', 'RandomElasticDeformation', 'RandomAffine', 'RandomFlip', 'OneOf']
 
 # %% ../nbs/03_vision_augment.ipynb 2
 from fastai.data.all import *
 from .vision_core import *
 import torchio as tio
 from monai.transforms import NormalizeIntensity as MonaiNormalizeIntensity
+from monai.transforms import SpatialPad as MonaiSpatialPad
 
 # %% ../nbs/03_vision_augment.ipynb 5
 class CustomDictTransform(ItemTransform):
@@ -92,6 +93,74 @@ class PadOrCrop(DisplayedTransform):
         return type(o)(self.pad_or_crop(o))
 
 # %% ../nbs/03_vision_augment.ipynb 9
+class SpatialPad(DisplayedTransform):
+    """Pad image to minimum size without cropping using MONAI's `SpatialPad`.
+    
+    This transform pads each dimension to AT LEAST the specified size. Dimensions
+    already larger than `spatial_size` are left unchanged (no cropping), making it
+    ideal for patch-based training where images must be at least as large as the
+    patch size.
+    
+    Uses zero padding (constant value 0) by default.
+    
+    Args:
+        spatial_size: Minimum size [x, y, z] for each dimension. Can be int (same
+            for all dims) or list. Usually set to patch_size for patch-based training.
+        mode: Padding mode. Default 'constant' means zero padding. Other options:
+            'edge', 'reflect', 'symmetric'.
+    
+    Example:
+        >>> # Image 512×512×48, need at least 256×256×96 for patches
+        >>> transform = SpatialPad(spatial_size=[256, 256, 96])
+        >>> # Result: 512×512×96 (x,y unchanged, z padded from 48→96)
+        
+    Note:
+        This differs from `PadOrCrop` which will crop dimensions larger than
+        the target size. Use `SpatialPad` when you want to preserve large
+        dimensions and only pad small ones.
+    """
+    
+    order = 0
+
+    def __init__(self, spatial_size, mode='constant'):
+        if not is_listy(spatial_size): 
+            spatial_size = [spatial_size, spatial_size, spatial_size]
+        self.spatial_size = spatial_size
+        self.mode = mode
+        # Create MONAI transform (zero padding by default)
+        self.spatial_pad = MonaiSpatialPad(
+            spatial_size=spatial_size,
+            mode=mode
+        )
+
+    @property
+    def tio_transform(self):
+        """Return TorchIO-compatible transform for patch-based workflows.
+        
+        This wraps MONAI's SpatialPad in a TorchIO SpatialTransform for
+        compatibility with TorchIO's Queue and Compose pipelines.
+        """
+        import torchio as tio
+        
+        class _TioSpatialPad(tio.SpatialTransform):
+            """TorchIO wrapper for MONAI's SpatialPad."""
+            def __init__(self, spatial_size, mode, **kwargs):
+                super().__init__(**kwargs)
+                self._monai_pad = MonaiSpatialPad(spatial_size=spatial_size, mode=mode)
+            
+            def apply_transform(self, subject):
+                for image in self.get_images(subject):
+                    data = image.data
+                    padded = self._monai_pad(data)
+                    image.set_data(padded)
+                return subject
+        
+        return _TioSpatialPad(spatial_size=self.spatial_size, mode=self.mode)
+
+    def encodes(self, o: (MedImage, MedMask)):
+        return type(o)(self.spatial_pad(o))
+
+# %% ../nbs/03_vision_augment.ipynb 10
 class ZNormalization(DisplayedTransform):
     """Apply TorchIO `ZNormalization`."""
 
@@ -136,7 +205,7 @@ class ZNormalization(DisplayedTransform):
     def encodes(self, o: MedMask):
         return o
 
-# %% ../nbs/03_vision_augment.ipynb 10
+# %% ../nbs/03_vision_augment.ipynb 11
 class RescaleIntensity(DisplayedTransform):
     """Apply TorchIO RescaleIntensity for robust intensity scaling.
     
@@ -165,7 +234,7 @@ class RescaleIntensity(DisplayedTransform):
     def encodes(self, o: MedMask):
         return o
 
-# %% ../nbs/03_vision_augment.ipynb 11
+# %% ../nbs/03_vision_augment.ipynb 12
 class NormalizeIntensity(DisplayedTransform):
     """Apply MONAI NormalizeIntensity.
     
@@ -203,7 +272,7 @@ class NormalizeIntensity(DisplayedTransform):
     def encodes(self, o: MedMask):
         return o
 
-# %% ../nbs/03_vision_augment.ipynb 12
+# %% ../nbs/03_vision_augment.ipynb 13
 class BraTSMaskConverter(DisplayedTransform):
     '''Convert BraTS masks.'''
 
@@ -215,7 +284,7 @@ class BraTSMaskConverter(DisplayedTransform):
         o = torch.where(o==4, 3., o)
         return MedMask.create(o)
 
-# %% ../nbs/03_vision_augment.ipynb 13
+# %% ../nbs/03_vision_augment.ipynb 14
 class BinaryConverter(DisplayedTransform):
     '''Convert to binary mask.'''
 
@@ -228,7 +297,7 @@ class BinaryConverter(DisplayedTransform):
         o = torch.where(o>0, 1., 0)
         return MedMask.create(o)
 
-# %% ../nbs/03_vision_augment.ipynb 14
+# %% ../nbs/03_vision_augment.ipynb 15
 class RandomGhosting(DisplayedTransform):
     """Apply TorchIO `RandomGhosting`."""
     
@@ -243,16 +312,12 @@ class RandomGhosting(DisplayedTransform):
         return self.add_ghosts
 
     def encodes(self, o: MedImage):
-        result = self.add_ghosts(o)
-        # Handle potential complex values from k-space operations
-        if result.is_complex():
-            result = torch.real(result)
-        return MedImage.create(result)
+        return MedImage.create(self.add_ghosts(o))
 
     def encodes(self, o: MedMask):
         return o
 
-# %% ../nbs/03_vision_augment.ipynb 15
+# %% ../nbs/03_vision_augment.ipynb 16
 class RandomSpike(DisplayedTransform):
     '''Apply TorchIO `RandomSpike`.'''
     
@@ -267,16 +332,12 @@ class RandomSpike(DisplayedTransform):
         return self.add_spikes
 
     def encodes(self, o: MedImage): 
-        result = self.add_spikes(o)
-        # Handle potential complex values from k-space operations
-        if result.is_complex():
-            result = torch.real(result)
-        return MedImage.create(result)
+        return MedImage.create(self.add_spikes(o))
         
     def encodes(self, o: MedMask):
         return o
 
-# %% ../nbs/03_vision_augment.ipynb 16
+# %% ../nbs/03_vision_augment.ipynb 17
 class RandomNoise(DisplayedTransform):
     '''Apply TorchIO `RandomNoise`.'''
 
@@ -296,7 +357,7 @@ class RandomNoise(DisplayedTransform):
     def encodes(self, o: MedMask):
         return o
 
-# %% ../nbs/03_vision_augment.ipynb 17
+# %% ../nbs/03_vision_augment.ipynb 18
 class RandomBiasField(DisplayedTransform):
     '''Apply TorchIO `RandomBiasField`.'''
 
@@ -316,7 +377,7 @@ class RandomBiasField(DisplayedTransform):
     def encodes(self, o: MedMask):
         return o
 
-# %% ../nbs/03_vision_augment.ipynb 18
+# %% ../nbs/03_vision_augment.ipynb 19
 class RandomBlur(DisplayedTransform):
     '''Apply TorchIO `RandomBlur`.'''
 
@@ -336,7 +397,7 @@ class RandomBlur(DisplayedTransform):
     def encodes(self, o: MedMask):
         return o
 
-# %% ../nbs/03_vision_augment.ipynb 19
+# %% ../nbs/03_vision_augment.ipynb 20
 class RandomGamma(DisplayedTransform):
     '''Apply TorchIO `RandomGamma`.'''
 
@@ -356,7 +417,7 @@ class RandomGamma(DisplayedTransform):
     def encodes(self, o: MedMask):
         return o
 
-# %% ../nbs/03_vision_augment.ipynb 20
+# %% ../nbs/03_vision_augment.ipynb 21
 class RandomIntensityScale(DisplayedTransform):
     """Randomly scale image intensities by a multiplicative factor.
 
@@ -388,7 +449,7 @@ class RandomIntensityScale(DisplayedTransform):
     def encodes(self, o: MedMask):
         return o
 
-# %% ../nbs/03_vision_augment.ipynb 21
+# %% ../nbs/03_vision_augment.ipynb 22
 class RandomMotion(DisplayedTransform):
     """Apply TorchIO `RandomMotion`."""
 
@@ -416,16 +477,249 @@ class RandomMotion(DisplayedTransform):
         return self.add_motion
 
     def encodes(self, o: MedImage):
-        result = self.add_motion(o)
-        # Handle potential complex values from k-space operations
-        if result.is_complex():
-            result = torch.real(result)
-        return MedImage.create(result)
+        return MedImage.create(self.add_motion(o))
 
     def encodes(self, o: MedMask):
         return o
 
 # %% ../nbs/03_vision_augment.ipynb 23
+def _create_ellipsoid_mask(shape, center, radii):
+    """Create a 3D ellipsoid mask.
+
+    Args:
+        shape: (D, H, W) shape of the volume
+        center: (z, y, x) center of ellipsoid
+        radii: (rz, ry, rx) radii along each axis
+
+    Returns:
+        Boolean mask where True = inside ellipsoid
+    """
+    z, y, x = torch.meshgrid(
+        torch.arange(shape[0]),
+        torch.arange(shape[1]),
+        torch.arange(shape[2]),
+        indexing='ij'
+    )
+    dist = ((z - center[0]) / radii[0]) ** 2 + \
+           ((y - center[1]) / radii[1]) ** 2 + \
+           ((x - center[2]) / radii[2]) ** 2
+    return dist <= 1.0
+
+# %% ../nbs/03_vision_augment.ipynb 24
+class _TioRandomCutout(tio.IntensityTransform):
+    """TorchIO-compatible RandomCutout for patch-based workflows.
+
+    When mask_only=True, cutouts only affect voxels where the mask is positive.
+    The mask should be available in the Subject as 'mask' key.
+    """
+
+    def __init__(self, holes=1, spatial_size=8, fill_value=None,
+                 max_holes=None, max_spatial_size=None, mask_only=True, p=0.2, **kwargs):
+        super().__init__(p=p, **kwargs)
+        self.holes = holes
+        self.spatial_size = spatial_size
+        self.fill_value = fill_value
+        self.max_holes = max_holes
+        self.max_spatial_size = max_spatial_size
+        self.mask_only = mask_only
+
+    def _apply_cutout(self, data, fill_val, mask_tensor=None):
+        """Apply spherical cutout(s) to a tensor.
+
+        Args:
+            data: Input tensor of shape (C, D, H, W)
+            fill_val: Value to fill cutout regions
+            mask_tensor: Optional mask tensor for mask-only cutouts
+
+        Returns:
+            Tensor with cutout applied
+        """
+        result = data.clone()
+        n_holes = torch.randint(self.holes, (self.max_holes or self.holes) + 1, (1,)).item()
+
+        spatial_shape = data.shape[1:]  # (D, H, W)
+        min_size = self.spatial_size if isinstance(self.spatial_size, int) else self.spatial_size[0]
+        max_size = self.max_spatial_size or self.spatial_size
+        max_size = max_size if isinstance(max_size, int) else max_size[0]
+
+        for _ in range(n_holes):
+            # Random size for this hole
+            size = torch.randint(min_size, max_size + 1, (3,))
+            radii = size.float() / 2
+
+            # Random center (ensure hole fits in volume)
+            center = [
+                torch.randint(int(radii[i].item()),
+                              max(spatial_shape[i] - int(radii[i].item()), int(radii[i].item()) + 1),
+                              (1,)).item()
+                for i in range(3)
+            ]
+
+            ellipsoid = _create_ellipsoid_mask(spatial_shape, center, radii)
+
+            if self.mask_only and mask_tensor is not None:
+                # INTERSECT with tumor mask - only affect tumor voxels
+                tumor_mask = mask_tensor[0] > 0
+                cutout_region = ellipsoid & tumor_mask
+            else:
+                cutout_region = ellipsoid
+
+            result[:, cutout_region] = fill_val
+
+        return result
+
+    def apply_transform(self, subject):
+        # Get mask if available for mask-only cutouts
+        mask_tensor = None
+        if self.mask_only and 'mask' in subject:
+            mask_tensor = subject['mask'].data
+
+            # Skip if mask is empty
+            if mask_tensor is not None and not (mask_tensor > 0).any():
+                return subject
+
+        for image in self.get_images(subject):
+            data = image.data
+            fill_val = self.fill_value if self.fill_value is not None else float(data.min())
+            result = self._apply_cutout(data, fill_val, mask_tensor)
+            image.set_data(result)
+        return subject
+
+# %% ../nbs/03_vision_augment.ipynb 25
+class RandomCutout(ItemTransform):
+    """Randomly erase spherical regions in 3D medical images with mask-aware placement.
+
+    Simulates post-operative surgical cavities by filling random ellipsoid
+    volumes with specified values. When mask_only=True (default), cutouts only
+    affect voxels inside the segmentation mask, ensuring no healthy tissue is modified.
+
+    Args:
+        holes: Minimum number of cutout regions. Default: 1.
+        max_holes: Maximum number of regions. Default: 3.
+        spatial_size: Minimum cutout diameter in voxels. Default: 8.
+        max_spatial_size: Maximum cutout diameter. Default: 16.
+        fill: Fill value - 'min', 'mean', 'random', or float. Default: 'min'.
+        mask_only: If True, cutouts only affect mask-positive voxels (tumor tissue).
+            If False, cutouts can affect any voxel (original behavior). Default: True.
+        p: Probability of applying transform. Default: 0.2.
+
+    Example:
+        >>> # Simulate post-op cavities only within tumor regions
+        >>> tfm = RandomCutout(holes=1, max_holes=2, spatial_size=10,
+        ...                    max_spatial_size=25, fill='min', mask_only=True, p=0.2)
+
+        >>> # Original behavior - cutouts anywhere in the volume
+        >>> tfm = RandomCutout(mask_only=False, p=0.2)
+    """
+
+    split_idx, order = 0, 1
+
+    def __init__(self, holes=1, max_holes=3, spatial_size=8,
+                 max_spatial_size=16, fill='min', mask_only=True, p=0.2):
+        self.holes = holes
+        self.max_holes = max_holes
+        self.spatial_size = spatial_size
+        self.max_spatial_size = max_spatial_size
+        self.fill = fill
+        self.mask_only = mask_only
+        self.p = p
+
+        self._tio_cutout = _TioRandomCutout(
+            holes=holes, spatial_size=spatial_size,
+            fill_value=None if isinstance(fill, str) else fill,
+            max_holes=max_holes, max_spatial_size=max_spatial_size,
+            mask_only=mask_only, p=p
+        )
+
+    @property
+    def tio_transform(self):
+        """Return TorchIO-compatible transform for patch-based workflows.
+
+        Note: For mask-aware cutouts in patch workflows, the mask must be
+        available in the TorchIO Subject as 'mask' key.
+        """
+        return self._tio_cutout
+
+    def _get_fill_value(self, tensor):
+        if self.fill == 'min': return float(tensor.min())
+        elif self.fill == 'mean': return float(tensor.mean())
+        elif self.fill == 'random':
+            return torch.empty(1).uniform_(float(tensor.min()), float(tensor.max())).item()
+        else: return self.fill
+
+    def encodes(self, x):
+        """Apply mask-aware cutout to image.
+
+        Args:
+            x: Tuple of (MedImage, target) where target is MedMask or TensorCategory
+
+        Returns:
+            Tuple of (transformed MedImage, unchanged target)
+        """
+        img, y_true = x
+
+        # Probability check
+        if torch.rand(1).item() > self.p:
+            return img, y_true
+
+        # Get mask data if available (as numpy for safe boolean operations)
+        mask_np = None
+        tumor_coords = None
+        if isinstance(y_true, MedMask):
+            mask_np = y_true.numpy()
+            # Get coordinates of tumor voxels for mask-aware center placement
+            if self.mask_only and (mask_np > 0).any():
+                tumor_coords = np.argwhere(mask_np[0] > 0)  # Shape: (N, 3) for z, y, x
+
+        # Skip cutout if mask_only=True but no mask or empty mask
+        if self.mask_only:
+            if mask_np is None or tumor_coords is None or len(tumor_coords) == 0:
+                return img, y_true
+
+        # Work with numpy array to avoid tensor subclass issues
+        result_np = img.numpy().copy()
+        spatial_shape = img.shape[1:]  # (D, H, W)
+        fill_val = self._get_fill_value(img)
+
+        n_holes = torch.randint(self.holes, self.max_holes + 1, (1,)).item()
+
+        min_size = self.spatial_size if isinstance(self.spatial_size, int) else self.spatial_size[0]
+        max_size = self.max_spatial_size or self.spatial_size
+        max_size = max_size if isinstance(max_size, int) else max_size[0]
+
+        for _ in range(n_holes):
+            # Random size for this hole
+            size = torch.randint(min_size, max_size + 1, (3,))
+            radii = size.float() / 2
+
+            if self.mask_only and tumor_coords is not None:
+                # Pick center from within tumor region to ensure intersection
+                idx = torch.randint(0, len(tumor_coords), (1,)).item()
+                center = tumor_coords[idx].tolist()  # [z, y, x]
+            else:
+                # Random center anywhere (ensure hole fits in volume)
+                center = [
+                    torch.randint(int(radii[i].item()),
+                                  max(spatial_shape[i] - int(radii[i].item()), int(radii[i].item()) + 1),
+                                  (1,)).item()
+                    for i in range(3)
+                ]
+
+            # Create ellipsoid mask (numpy)
+            ellipsoid = _create_ellipsoid_mask(spatial_shape, center, radii).numpy()
+
+            if self.mask_only and mask_np is not None:
+                # INTERSECT with tumor mask - only affect tumor voxels
+                tumor_mask = mask_np[0] > 0
+                cutout_region = ellipsoid & tumor_mask
+            else:
+                cutout_region = ellipsoid
+
+            result_np[:, cutout_region] = fill_val
+
+        return MedImage.create(torch.from_numpy(result_np)), y_true
+
+# %% ../nbs/03_vision_augment.ipynb 27
 class RandomElasticDeformation(CustomDictTransform):
     """Apply TorchIO `RandomElasticDeformation`."""
 
@@ -438,7 +732,7 @@ class RandomElasticDeformation(CustomDictTransform):
             image_interpolation=image_interpolation,
             p=p))
 
-# %% ../nbs/03_vision_augment.ipynb 24
+# %% ../nbs/03_vision_augment.ipynb 28
 class RandomAffine(CustomDictTransform):
     """Apply TorchIO `RandomAffine`."""
 
@@ -454,14 +748,14 @@ class RandomAffine(CustomDictTransform):
             default_pad_value=default_pad_value,
             p=p))
 
-# %% ../nbs/03_vision_augment.ipynb 25
+# %% ../nbs/03_vision_augment.ipynb 29
 class RandomFlip(CustomDictTransform):
     """Apply TorchIO `RandomFlip`."""
 
     def __init__(self, axes='LR', p=0.5):
         super().__init__(tio.RandomFlip(axes=axes, flip_probability=p))
 
-# %% ../nbs/03_vision_augment.ipynb 26
+# %% ../nbs/03_vision_augment.ipynb 30
 class OneOf(CustomDictTransform):
     """Apply only one of the given transforms using TorchIO `OneOf`."""