fastMONAI 0.5.3__py3-none-any.whl → 0.6.0__py3-none-any.whl

fastMONAI/vision_augmentation.py

@@ -3,8 +3,8 @@
 # %% auto 0
 __all__ = ['CustomDictTransform', 'do_pad_or_crop', 'PadOrCrop', 'ZNormalization', 'RescaleIntensity', 'NormalizeIntensity',
            'BraTSMaskConverter', 'BinaryConverter', 'RandomGhosting', 'RandomSpike', 'RandomNoise', 'RandomBiasField',
-           'RandomBlur', 'RandomGamma', 'RandomMotion', 'RandomElasticDeformation', 'RandomAffine', 'RandomFlip',
-           'OneOf']
+           'RandomBlur', 'RandomGamma', 'RandomIntensityScale', 'RandomMotion', 'RandomCutout',
+           'RandomElasticDeformation', 'RandomAffine', 'RandomFlip', 'OneOf']
 
 # %% ../nbs/03_vision_augment.ipynb 2
 from fastai.data.all import *
@@ -14,10 +14,10 @@ from monai.transforms import NormalizeIntensity as MonaiNormalizeIntensity
 
 # %% ../nbs/03_vision_augment.ipynb 5
 class CustomDictTransform(ItemTransform):
-    """A class that serves as a wrapper to perform an identical transformation on both 
+    """A class that serves as a wrapper to perform an identical transformation on both
     the image and the target (if it's a mask).
     """
-    
+
     split_idx = 0  # Only perform transformations on training data. Use TTA() for transformations on validation data.
 
     def __init__(self, aug):
@@ -28,31 +28,42 @@ class CustomDictTransform(ItemTransform):
         """
         self.aug = aug
 
+    @property
+    def tio_transform(self):
+        """Return the underlying TorchIO transform.
+
+        This property enables using fastMONAI wrappers in patch-based workflows
+        where raw TorchIO transforms are needed for tio.Compose().
+        """
+        return self.aug
+
     def encodes(self, x):
         """
-        Applies the stored transformation to an image, and the same random transformation 
+        Applies the stored transformation to an image, and the same random transformation
         to the target if it is a mask. If the target is not a mask, it returns the target as is.
 
         Args:
-            x (Tuple[MedImage, Union[MedMask, TensorCategory]]): A tuple containing the 
+            x (Tuple[MedImage, Union[MedMask, TensorCategory]]): A tuple containing the
             image and the target.
 
         Returns:
-            Tuple[MedImage, Union[MedMask, TensorCategory]]: The transformed image and target. 
-            If the target is a mask, it's transformed identically to the image. If the target 
+            Tuple[MedImage, Union[MedMask, TensorCategory]]: The transformed image and target.
+            If the target is a mask, it's transformed identically to the image. If the target
             is not a mask, the original target is returned.
         """
         img, y_true = x
+        
+        # Use identity affine if MedImage.affine_matrix is not set
+        affine = MedImage.affine_matrix if MedImage.affine_matrix is not None else np.eye(4)
 
         if isinstance(y_true, (MedMask)):
-            aug = self.aug(tio.Subject(img=tio.ScalarImage(tensor=img, affine=MedImage.affine_matrix), 
-                                        mask=tio.LabelMap(tensor=y_true, affine=MedImage.affine_matrix)))
+            aug = self.aug(tio.Subject(img=tio.ScalarImage(tensor=img, affine=affine),
+                                        mask=tio.LabelMap(tensor=y_true, affine=affine)))
             return MedImage.create(aug['img'].data), MedMask.create(aug['mask'].data)
 
         aug = self.aug(tio.Subject(img=tio.ScalarImage(tensor=img)))
         return MedImage.create(aug['img'].data), y_true
 
-
 # %% ../nbs/03_vision_augment.ipynb 7
 def do_pad_or_crop(o, target_shape, padding_mode, mask_name, dtype=torch.Tensor):
     #TODO:refactorize
@@ -72,6 +83,11 @@ class PadOrCrop(DisplayedTransform):
                                     padding_mode=padding_mode, 
                                     mask_name=mask_name)
 
+    @property
+    def tio_transform(self):
+        """Return the underlying TorchIO transform."""
+        return self.pad_or_crop
+
     def encodes(self, o: (MedImage, MedMask)):
         return type(o)(self.pad_or_crop(o))
 
@@ -85,6 +101,11 @@ class ZNormalization(DisplayedTransform):
         self.z_normalization = tio.ZNormalization(masking_method=masking_method)
         self.channel_wise = channel_wise
 
+    @property
+    def tio_transform(self):
+        """Return the underlying TorchIO transform."""
+        return self.z_normalization
+
     def encodes(self, o: MedImage):
         try:
             if self.channel_wise:
@@ -132,7 +153,12 @@ class RescaleIntensity(DisplayedTransform):
     
     def __init__(self, out_min_max: tuple[float, float], in_min_max: tuple[float, float]):
         self.rescale = tio.RescaleIntensity(out_min_max=out_min_max, in_min_max=in_min_max)
-    
+
+    @property
+    def tio_transform(self):
+        """Return the underlying TorchIO transform."""
+        return self.rescale
+
     def encodes(self, o: MedImage):
         return MedImage.create(self.rescale(o))
         
@@ -211,8 +237,17 @@ class RandomGhosting(DisplayedTransform):
     def __init__(self, intensity=(0.5, 1), p=0.5):
         self.add_ghosts = tio.RandomGhosting(intensity=intensity, p=p)
 
+    @property
+    def tio_transform(self):
+        """Return the underlying TorchIO transform."""
+        return self.add_ghosts
+
     def encodes(self, o: MedImage):
-        return MedImage.create(self.add_ghosts(o))
+        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)
 
     def encodes(self, o: MedMask):
         return o
@@ -221,26 +256,40 @@ class RandomGhosting(DisplayedTransform):
 class RandomSpike(DisplayedTransform):
     '''Apply TorchIO `RandomSpike`.'''
     
-    split_idx,order=0,1
+    split_idx, order = 0, 1
 
     def __init__(self, num_spikes=1, intensity=(1, 3), p=0.5):
         self.add_spikes = tio.RandomSpike(num_spikes=num_spikes, intensity=intensity, p=p)
 
-    def encodes(self, o:MedImage): 
-        return MedImage.create(self.add_spikes(o))
+    @property
+    def tio_transform(self):
+        """Return the underlying TorchIO transform."""
+        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)
         
-    def encodes(self, o:MedMask):
+    def encodes(self, o: MedMask):
         return o
 
 # %% ../nbs/03_vision_augment.ipynb 16
 class RandomNoise(DisplayedTransform):
     '''Apply TorchIO `RandomNoise`.'''
 
-    split_idx,order=0,1
+    split_idx, order = 0, 1
 
     def __init__(self, mean=0, std=(0, 0.25), p=0.5):
         self.add_noise = tio.RandomNoise(mean=mean, std=std, p=p)
 
+    @property
+    def tio_transform(self):
+        """Return the underlying TorchIO transform."""
+        return self.add_noise
+
     def encodes(self, o: MedImage): 
         return MedImage.create(self.add_noise(o))
     
@@ -251,11 +300,16 @@ class RandomNoise(DisplayedTransform):
 class RandomBiasField(DisplayedTransform):
     '''Apply TorchIO `RandomBiasField`.'''
 
-    split_idx,order=0,1
+    split_idx, order = 0, 1
 
     def __init__(self, coefficients=0.5, order=3, p=0.5):
         self.add_biasfield = tio.RandomBiasField(coefficients=coefficients, order=order, p=p)
 
+    @property
+    def tio_transform(self):
+        """Return the underlying TorchIO transform."""
+        return self.add_biasfield
+
     def encodes(self, o: MedImage): 
         return MedImage.create(self.add_biasfield(o))
         
@@ -264,13 +318,18 @@ class RandomBiasField(DisplayedTransform):
 
 # %% ../nbs/03_vision_augment.ipynb 18
 class RandomBlur(DisplayedTransform):
-    '''Apply TorchIO `RandomBiasField`.'''
+    '''Apply TorchIO `RandomBlur`.'''
 
-    split_idx,order=0,1
+    split_idx, order = 0, 1
 
     def __init__(self, std=(0, 2), p=0.5):
         self.add_blur = tio.RandomBlur(std=std, p=p)
-        
+
+    @property
+    def tio_transform(self):
+        """Return the underlying TorchIO transform."""
+        return self.add_blur
+
     def encodes(self, o: MedImage): 
         return MedImage.create(self.add_blur(o))
     
@@ -281,12 +340,16 @@ class RandomBlur(DisplayedTransform):
 class RandomGamma(DisplayedTransform):
     '''Apply TorchIO `RandomGamma`.'''
 
-
-    split_idx,order=0,1
+    split_idx, order = 0, 1
 
     def __init__(self, log_gamma=(-0.3, 0.3), p=0.5):
         self.add_gamma = tio.RandomGamma(log_gamma=log_gamma, p=p)
 
+    @property
+    def tio_transform(self):
+        """Return the underlying TorchIO transform."""
+        return self.add_gamma
+
     def encodes(self, o: MedImage): 
         return MedImage.create(self.add_gamma(o))
     
@@ -294,6 +357,38 @@ class RandomGamma(DisplayedTransform):
         return o
 
 # %% ../nbs/03_vision_augment.ipynb 20
+class RandomIntensityScale(DisplayedTransform):
+    """Randomly scale image intensities by a multiplicative factor.
+
+    Useful for domain generalization across different acquisition protocols
+    with varying intensity ranges.
+
+    Args:
+        scale_range (tuple[float, float]): Range of scale factors (min, max).
+            Values > 1 increase intensity, < 1 decrease intensity.
+        p (float): Probability of applying the transform (default: 0.5)
+
+    Example:
+        # Scale intensities randomly between 0.5x and 2.0x
+        transform = RandomIntensityScale(scale_range=(0.5, 2.0), p=0.3)
+    """
+
+    split_idx, order = 0, 1
+
+    def __init__(self, scale_range: tuple[float, float] = (0.5, 2.0), p: float = 0.5):
+        self.scale_range = scale_range
+        self.p = p
+
+    def encodes(self, o: MedImage):
+        if torch.rand(1).item() > self.p:
+            return o
+        scale = torch.empty(1).uniform_(self.scale_range[0], self.scale_range[1]).item()
+        return MedImage.create(o * scale)
+
+    def encodes(self, o: MedMask):
+        return o
+
+# %% ../nbs/03_vision_augment.ipynb 21
 class RandomMotion(DisplayedTransform):
     """Apply TorchIO `RandomMotion`."""
 
@@ -315,13 +410,154 @@ class RandomMotion(DisplayedTransform):
             p=p
         )
 
+    @property
+    def tio_transform(self):
+        """Return the underlying TorchIO transform."""
+        return self.add_motion
+
     def encodes(self, o: MedImage):
-        return MedImage.create(self.add_motion(o))
+        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)
 
     def encodes(self, o: MedMask):
         return o
 
 # %% ../nbs/03_vision_augment.ipynb 22
+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 23
+class _TioRandomCutout(tio.IntensityTransform):
+    """TorchIO-compatible RandomCutout for patch-based workflows."""
+
+    def __init__(self, holes=1, spatial_size=8, fill_value=None,
+                 max_holes=None, max_spatial_size=None, 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
+
+    def _apply_cutout(self, data, fill_val):
+        """Apply spherical cutout(s) to a tensor."""
+        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)
+            ]
+
+            mask = _create_ellipsoid_mask(spatial_shape, center, radii)
+            result[:, mask] = fill_val
+
+        return result
+
+    def apply_transform(self, 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)
+            image.set_data(result)
+        return subject
+
+# %% ../nbs/03_vision_augment.ipynb 24
+class RandomCutout(DisplayedTransform):
+    """Randomly erase spherical regions in 3D medical images.
+
+    Simulates post-operative surgical cavities by filling random ellipsoid
+    volumes with specified values. Useful for training on pre-op images
+    to generalize to post-op scans.
+
+    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'.
+        p: Probability of applying transform. Default: 0.2.
+
+    Example:
+        >>> # Simulate post-op cavities with dark spherical voids
+        >>> tfm = RandomCutout(holes=1, max_holes=2, spatial_size=10,
+        ...                    max_spatial_size=25, fill='min', 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', 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.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, p=p
+        )
+
+    @property
+    def tio_transform(self):
+        """Return TorchIO-compatible transform for patch-based workflows."""
+        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, o: MedImage):
+        if torch.rand(1).item() > self.p: return o
+        fill_val = self._get_fill_value(o)
+        result = self._tio_cutout._apply_cutout(o.clone(), fill_val)
+        return MedImage.create(result)
+
+    def encodes(self, o: MedMask):
+        return o
+
+# %% ../nbs/03_vision_augment.ipynb 26
 class RandomElasticDeformation(CustomDictTransform):
     """Apply TorchIO `RandomElasticDeformation`."""
 
@@ -334,7 +570,7 @@ class RandomElasticDeformation(CustomDictTransform):
             image_interpolation=image_interpolation,
             p=p))
 
-# %% ../nbs/03_vision_augment.ipynb 23
+# %% ../nbs/03_vision_augment.ipynb 27
 class RandomAffine(CustomDictTransform):
     """Apply TorchIO `RandomAffine`."""
 
@@ -350,14 +586,14 @@ class RandomAffine(CustomDictTransform):
             default_pad_value=default_pad_value,
             p=p))
 
-# %% ../nbs/03_vision_augment.ipynb 24
+# %% ../nbs/03_vision_augment.ipynb 28
 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 25
+# %% ../nbs/03_vision_augment.ipynb 29
 class OneOf(CustomDictTransform):
     """Apply only one of the given transforms using TorchIO `OneOf`."""