fastMONAI 0.3.1__py3-none-any.whl → 0.3.3__py3-none-any.whl

fastMONAI/vision_augmentation.py

@@ -12,75 +12,88 @@ import torchio as tio
 
 # %% ../nbs/03_vision_augment.ipynb 5
 class CustomDictTransform(ItemTransform):
-    '''Wrapper to perform an identical transformation on both image and target (if it is a mask) during training.'''
+    """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
-    def __init__(self, aug): self.aug = aug
+    split_idx = 0  # Only perform transformations on training data. Use TTA() for transformations on validation data.
+
+    def __init__(self, aug):
+        """Constructs CustomDictTransform object.
+
+        Args:
+            aug (Callable): Function to apply augmentation on the image.
+        """
+        self.aug = aug
 
     def encodes(self, x):
-        '''Apply transformation to an image, and the same random transformation to the target if it is a mask.
+        """
+        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: Contains image and target.
+            x (Tuple[MedImage, Union[MedMask, TensorCategory]]): A tuple containing the 
+            image and the target.
 
         Returns:
-            MedImage: Transformed image data.
-            (MedMask, TensorCategory, ...todo): If the target is a mask, then return a transformed mask data. Otherwise, return target value.
-        '''
-
+            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
 
         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=MedImage.affine_matrix), 
+                                        mask=tio.LabelMap(tensor=y_true, affine=MedImage.affine_matrix)))
             return MedImage.create(aug['img'].data), MedMask.create(aug['mask'].data)
-        else:
-            aug = self.aug(tio.Subject(img=tio.ScalarImage(tensor=img)))
-            return MedImage.create(aug['img'].data), y_true
 
-# %% ../nbs/03_vision_augment.ipynb 8
-def do_pad_or_crop(o, target_shape, padding_mode, mask_name, dtype=torch.Tensor):
+        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
     pad_or_crop = tio.CropOrPad(target_shape=target_shape, padding_mode=padding_mode, mask_name=mask_name)
     return dtype(pad_or_crop(o))
 
-# %% ../nbs/03_vision_augment.ipynb 9
+# %% ../nbs/03_vision_augment.ipynb 8
 class PadOrCrop(DisplayedTransform):
-    '''Resize image using TorchIO `CropOrPad`.'''
+    """Resize image using TorchIO `CropOrPad`."""
+    
+    order = 0
 
-    order=0
     def __init__(self, size, padding_mode=0, mask_name=None):
-        if not is_listy(size): size=[size,size,size]
-        self.size, self.padding_mode, self.mask_name = size, padding_mode, mask_name
+        if not is_listy(size): 
+            size = [size, size, size]
+        self.pad_or_crop = tio.CropOrPad(target_shape=size,
+                                    padding_mode=padding_mode, 
+                                    mask_name=mask_name)
 
-    def encodes(self, o:(MedImage, MedMask)):
-        return do_pad_or_crop(o,target_shape=self.size, padding_mode=self.padding_mode, mask_name=self.mask_name, dtype=type(o))
+    def encodes(self, o: (MedImage, MedMask)):
+        return type(o)(self.pad_or_crop(o))
 
-# %% ../nbs/03_vision_augment.ipynb 11
-def _do_z_normalization(o, masking_method, channel_wise):
-
-    z_normalization = tio.ZNormalization(masking_method=masking_method)
-    normalized_tensor = torch.zeros(o.shape)
+# %% ../nbs/03_vision_augment.ipynb 9
+class ZNormalization(DisplayedTransform):
+    """Apply TorchIO `ZNormalization`."""
 
-    if channel_wise:
-        for idx, c in enumerate(o): 
-            normalized_tensor[idx] = z_normalization(c[None])[0]
-            
-    else: normalized_tensor = z_normalization(o)
+    order = 0
 
-    return normalized_tensor
+    def __init__(self, masking_method=None, channel_wise=True):
+        self.z_normalization = tio.ZNormalization(masking_method=masking_method)
+        self.channel_wise = channel_wise
 
-# %% ../nbs/03_vision_augment.ipynb 12
-class ZNormalization(DisplayedTransform):
-    '''Apply TorchIO `ZNormalization`.'''
+    def encodes(self, o: MedImage):
+        if self.channel_wise:
+            o = torch.stack([self.z_normalization(c[None])[0] for c in o])
+        else: o = self.z_normalization(o) 
 
-    order=0
-    def __init__(self, masking_method=None, channel_wise=True):
-        self.masking_method, self.channel_wise = masking_method, channel_wise
+        return MedImage.create(o)
 
-    def encodes(self, o:(MedImage)): return MedImage.create(_do_z_normalization(o, self.masking_method, self.channel_wise))
-    def encodes(self, o:(MedMask)):return o
+    def encodes(self, o: MedMask):
+        return o
 
-# %% ../nbs/03_vision_augment.ipynb 14
+# %% ../nbs/03_vision_augment.ipynb 10
 class BraTSMaskConverter(DisplayedTransform):
     '''Convert BraTS masks.'''
 
@@ -92,115 +105,95 @@ class BraTSMaskConverter(DisplayedTransform):
         o = torch.where(o==4, 3., o)
         return MedMask.create(o)
 
-# %% ../nbs/03_vision_augment.ipynb 16
+# %% ../nbs/03_vision_augment.ipynb 11
 class BinaryConverter(DisplayedTransform):
     '''Convert to binary mask.'''
 
     order=1
 
-    def encodes(self, o:(MedImage)): return o
+    def encodes(self, o: MedImage): 
+        return o
 
-    def encodes(self, o:(MedMask)):
+    def encodes(self, o: MedMask):
         o = torch.where(o>0, 1., 0)
         return MedMask.create(o)
 
-# %% ../nbs/03_vision_augment.ipynb 18
-def _do_rand_ghosting(o, intensity, p):
-    
-    add_ghosts = tio.RandomGhosting(intensity=intensity, p=p)
-    return add_ghosts(o)
-
-# %% ../nbs/03_vision_augment.ipynb 19
+# %% ../nbs/03_vision_augment.ipynb 12
 class RandomGhosting(DisplayedTransform):
-    '''Apply TorchIO `RandomGhosting`.'''
-
-    split_idx,order=0,1
-
-    def __init__(self, intensity =(0.5, 1), p=0.5):
-        self.intensity, self.p  = intensity, p
+    """Apply TorchIO `RandomGhosting`."""
+    
+    split_idx, order = 0, 1
 
-    def encodes(self, o:(MedImage)): return MedImage.create(_do_rand_ghosting(o, self.intensity, self.p))
-    def encodes(self, o:(MedMask)):return o
+    def __init__(self, intensity=(0.5, 1), p=0.5):
+        self.add_ghosts = tio.RandomGhosting(intensity=intensity, p=p)
 
-# %% ../nbs/03_vision_augment.ipynb 21
-def _do_rand_spike(o, num_spikes, intensity, p):
+    def encodes(self, o: MedImage):
+        return MedImage.create(self.add_ghosts(o))
 
-    add_spikes = tio.RandomSpike(num_spikes=num_spikes, intensity=intensity, p=p)
-    return add_spikes(o) #return torch tensor
+    def encodes(self, o: MedMask):
+        return o
 
-# %% ../nbs/03_vision_augment.ipynb 22
+# %% ../nbs/03_vision_augment.ipynb 13
 class RandomSpike(DisplayedTransform):
     '''Apply TorchIO `RandomSpike`.'''
     
     split_idx,order=0,1
 
     def __init__(self, num_spikes=1, intensity=(1, 3), p=0.5):
-        self.num_spikes, self.intensity, self.p  = num_spikes, intensity, p
+        self.add_spikes = tio.RandomSpike(num_spikes=num_spikes, intensity=intensity, p=p)
 
-    def encodes(self, o:(MedImage)): return MedImage.create(_do_rand_spike(o, self.num_spikes, self.intensity, self.p))
-    def encodes(self, o:(MedMask)):return o
+    def encodes(self, o:MedImage): 
+        return MedImage.create(self.add_spikes(o))
+        
+    def encodes(self, o:MedMask):
+        return o
 
-# %% ../nbs/03_vision_augment.ipynb 24
-def _do_rand_noise(o, mean, std, p):
-
-    add_noise = tio.RandomNoise(mean=mean, std=std, p=p)
-    return add_noise(o) #return torch tensor
-
-# %% ../nbs/03_vision_augment.ipynb 25
+# %% ../nbs/03_vision_augment.ipynb 14
 class RandomNoise(DisplayedTransform):
     '''Apply TorchIO `RandomNoise`.'''
 
     split_idx,order=0,1
 
     def __init__(self, mean=0, std=(0, 0.25), p=0.5):
-        self.mean, self.std, self.p  = mean, std, p
-
-    def encodes(self, o:(MedImage)): return MedImage.create(_do_rand_noise(o, mean=self.mean, std=self.std, p=self.p))
-    def encodes(self, o:(MedMask)):return o
+        self.add_noise = tio.RandomNoise(mean=mean, std=std, p=p)
 
-# %% ../nbs/03_vision_augment.ipynb 27
-def _do_rand_biasfield(o, coefficients, order, p):
-
-    add_biasfield = tio.RandomBiasField(coefficients=coefficients, order=order, p=p)
-    return add_biasfield(o) #return torch tensor
+    def encodes(self, o: MedImage): 
+        return MedImage.create(self.add_noise(o))
+    
+    def encodes(self, o: MedMask):
+        return o
 
-# %% ../nbs/03_vision_augment.ipynb 28
+# %% ../nbs/03_vision_augment.ipynb 15
 class RandomBiasField(DisplayedTransform):
     '''Apply TorchIO `RandomBiasField`.'''
 
     split_idx,order=0,1
 
     def __init__(self, coefficients=0.5, order=3, p=0.5):
-        self.coefficients, self.order, self.p  = coefficients, order, p
-
-    def encodes(self, o:(MedImage)): return MedImage.create(_do_rand_biasfield(o, coefficients=self.coefficients, order=self.order, p=self.p))
-    def encodes(self, o:(MedMask)):return o
-
-# %% ../nbs/03_vision_augment.ipynb 30
-def _do_rand_blur(o, std, p):
+        self.add_biasfield = tio.RandomBiasField(coefficients=coefficients, order=order, p=p)
 
-    add_blur = tio.RandomBlur(std=std, p=p)
-    return add_blur(o) 
+    def encodes(self, o: MedImage): 
+        return MedImage.create(self.add_biasfield(o))
+        
+    def encodes(self, o: MedMask):
+        return o
 
-# %% ../nbs/03_vision_augment.ipynb 31
+# %% ../nbs/03_vision_augment.ipynb 16
 class RandomBlur(DisplayedTransform):
     '''Apply TorchIO `RandomBiasField`.'''
 
     split_idx,order=0,1
 
     def __init__(self, std=(0, 2), p=0.5):
-        self.std, self.p  = std, p
-
-    def encodes(self, o:(MedImage)): return MedImage.create(_do_rand_blur(o, std=self.std, p=self.p))
-    def encodes(self, o:(MedMask)):return o
-
-# %% ../nbs/03_vision_augment.ipynb 33
-def _do_rand_gamma(o, log_gamma, p):
-
-    add_gamma = tio.RandomGamma(log_gamma=log_gamma, p=p)
-    return add_gamma(o) 
+        self.add_blur = tio.RandomBlur(std=std, p=p)
+        
+    def encodes(self, o: MedImage): 
+        return MedImage.create(self.add_blur(o))
+    
+    def encodes(self, o: MedMask):
+        return o
 
-# %% ../nbs/03_vision_augment.ipynb 34
+# %% ../nbs/03_vision_augment.ipynb 17
 class RandomGamma(DisplayedTransform):
     '''Apply TorchIO `RandomGamma`.'''
 
@@ -208,53 +201,81 @@ class RandomGamma(DisplayedTransform):
     split_idx,order=0,1
 
     def __init__(self, log_gamma=(-0.3, 0.3), p=0.5):
-        self.log_gamma, self.p  = log_gamma, p
-
-    def encodes(self, o:(MedImage)): return MedImage.create(_do_rand_gamma(o, log_gamma=self.log_gamma, p=self.p))
-    def encodes(self, o:(MedMask)):return o
-
-# %% ../nbs/03_vision_augment.ipynb 36
-def _do_rand_motion(o, degrees, translation, num_transforms, image_interpolation, p):
+        self.add_gamma = tio.RandomGamma(log_gamma=log_gamma, p=p)
 
-    add_motion = tio.RandomMotion(degrees=degrees, translation=translation, num_transforms=num_transforms, image_interpolation=image_interpolation, p=p)
-    return add_motion(o) #return torch tensor
+    def encodes(self, o: MedImage): 
+        return MedImage.create(self.add_gamma(o))
+    
+    def encodes(self, o: MedMask):
+        return o
 
-# %% ../nbs/03_vision_augment.ipynb 37
+# %% ../nbs/03_vision_augment.ipynb 18
 class RandomMotion(DisplayedTransform):
-    '''Apply TorchIO `RandomMotion`.'''
-
-    split_idx,order=0,1
-
-    def __init__(self, degrees=10, translation=10, num_transforms=2, image_interpolation='linear', p=0.5):
-        self.degrees,self.translation, self.num_transforms, self.image_interpolation, self.p = degrees,translation, num_transforms, image_interpolation, p
-
-    def encodes(self, o:(MedImage)): return MedImage.create(_do_rand_motion(o, degrees=self.degrees,translation=self.translation, num_transforms=self.num_transforms, image_interpolation=self.image_interpolation, p=self.p))
-    def encodes(self, o:(MedMask)):return o
-
-# %% ../nbs/03_vision_augment.ipynb 40
+    """Apply TorchIO `RandomMotion`."""
+
+    split_idx, order = 0, 1
+
+    def __init__(
+            self, 
+            degrees=10, 
+            translation=10, 
+            num_transforms=2, 
+            image_interpolation='linear', 
+            p=0.5
+        ):
+        self.add_motion = tio.RandomMotion(
+            degrees=degrees, 
+            translation=translation, 
+            num_transforms=num_transforms, 
+            image_interpolation=image_interpolation, 
+            p=p
+        )
+
+    def encodes(self, o: MedImage):
+        return MedImage.create(self.add_motion(o))
+
+    def encodes(self, o: MedMask):
+        return o
+
+# %% ../nbs/03_vision_augment.ipynb 20
 class RandomElasticDeformation(CustomDictTransform):
-    '''Apply TorchIO `RandomElasticDeformation`.'''
+    """Apply TorchIO `RandomElasticDeformation`."""
 
-    def __init__(self,num_control_points=7, max_displacement=7.5, image_interpolation='linear', p=0.5): 
-        super().__init__(tio.RandomElasticDeformation(num_control_points=num_control_points, max_displacement=max_displacement, image_interpolation=image_interpolation, p=p))
+    def __init__(self, num_control_points=7, max_displacement=7.5,
+                 image_interpolation='linear', p=0.5):
+                     
+        super().__init__(tio.RandomElasticDeformation(
+            num_control_points=num_control_points,
+            max_displacement=max_displacement,
+            image_interpolation=image_interpolation,
+            p=p))
 
-# %% ../nbs/03_vision_augment.ipynb 42
+# %% ../nbs/03_vision_augment.ipynb 21
 class RandomAffine(CustomDictTransform):
-    '''Apply TorchIO `RandomAffine`.'''
-
-    def __init__(self, scales=0, degrees=10, translation=0, isotropic=False, image_interpolation='linear', default_pad_value=0., p=0.5): 
-        super().__init__(tio.RandomAffine(scales=scales, degrees=degrees, translation=translation, isotropic=isotropic, image_interpolation=image_interpolation, default_pad_value=default_pad_value, p=p))
+    """Apply TorchIO `RandomAffine`."""
+
+    def __init__(self, scales=0, degrees=10, translation=0, isotropic=False,
+                 image_interpolation='linear', default_pad_value=0., p=0.5):
+                     
+        super().__init__(tio.RandomAffine(
+            scales=scales,
+            degrees=degrees,
+            translation=translation,
+            isotropic=isotropic,
+            image_interpolation=image_interpolation,
+            default_pad_value=default_pad_value,
+            p=p))
 
-# %% ../nbs/03_vision_augment.ipynb 44
+# %% ../nbs/03_vision_augment.ipynb 22
 class RandomFlip(CustomDictTransform):
-    '''Apply TorchIO `RandomFlip`.'''
+    """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 46
+# %% ../nbs/03_vision_augment.ipynb 23
 class OneOf(CustomDictTransform):
-    '''Apply only one of the given transforms using TorchIO `OneOf`.'''
+    """Apply only one of the given transforms using TorchIO `OneOf`."""
 
     def __init__(self, transform_dict, p=1):
         super().__init__(tio.OneOf(transform_dict, p=p))

fastMONAI/vision_core.py

@@ -10,7 +10,8 @@ from torchio import ScalarImage, LabelMap, ToCanonical, Resample
 
 # %% ../nbs/01_vision_core.ipynb 5
 def _preprocess(obj, reorder, resample):
-    """Preprocesses the given object.
+    """
+    Preprocesses the given object.
 
     Args:
         obj: The object to preprocess.
@@ -83,12 +84,8 @@ def _multi_channel(image_paths: list, reorder: bool, resample: list, dtype, only
 
 
 # %% ../nbs/01_vision_core.ipynb 8
-def med_img_reader(
-        file_path: (str, Path),
-        dtype=torch.Tensor,
-        reorder: bool = False,
-        resample: list = None,
-        only_tensor: bool = True
+def med_img_reader(file_path: (str, Path), dtype=torch.Tensor, reorder: bool = False,
+                   resample: list = None, only_tensor: bool = True
 ):
     """Loads and preprocesses a medical image.
 
@@ -120,32 +117,36 @@ def med_img_reader(
 
 # %% ../nbs/01_vision_core.ipynb 10
 class MetaResolver(type(torch.Tensor), metaclass=BypassNewMeta):
-    '''A class to bypass metaclass conflict:
+    """
+    A class to bypass metaclass conflict:
     https://pytorch-geometric.readthedocs.io/en/latest/_modules/torch_geometric/data/batch.html
-    '''
+    """
     pass
 
 # %% ../nbs/01_vision_core.ipynb 11
-class MedBase(torch.Tensor, metaclass=MetaResolver): 
-    '''A class that represents an image object. Metaclass casts x to this class if it is of type cls._bypass_type.'''
-
-    _bypass_type=torch.Tensor
+class MedBase(torch.Tensor, metaclass=MetaResolver):
+    """A class that represents an image object.
+    Metaclass casts `x` to this class if it is of type `cls._bypass_type`."""
+    
+    _bypass_type = torch.Tensor
     _show_args = {'cmap':'gray'}
     resample, reorder = None, False
     affine_matrix = None
 
-
     @classmethod
-    def create(cls, fn: (Path, str, torch.Tensor), **kwargs):
+    def create(cls, fn: (Path, str, torch.Tensor), **kwargs) -> torch.Tensor:
         """
-        Open a medical image and cast to MedBase object. If it is a torch.Tensor, cast to MedBase object.
+        Opens a medical image and casts it to MedBase object.
+        If `fn` is a torch.Tensor, it's cast to MedBase object.
 
         Args:
-            fn: Image path or a 4D torch.Tensor.
-            kwargs: Additional parameters.
+            fn : (Path, str, torch.Tensor)
+                Image path or a 4D torch.Tensor.
+            kwargs : dict
+                Additional parameters for the medical image reader.
 
         Returns:
-            A 4D tensor as MedBase object.
+            torch.Tensor : A 4D tensor as a MedBase object.
         """
         if isinstance(fn, torch.Tensor):
             return cls(fn)
@@ -155,18 +156,32 @@ class MedBase(torch.Tensor, metaclass=MetaResolver):
     @classmethod
     def item_preprocessing(cls, resample: (list, int, tuple), reorder: bool):
         """
-        Change the values for the class variables `resample` and `reorder`.
+        Changes the values for the class variables `resample` and `reorder`.
 
         Args:
-            resample: A list with voxel spacing.
-            reorder: Whether to reorder the data to be closest to canonical (RAS+) orientation.
+            resample : (list, int, tuple)
+                A list with voxel spacing.
+            reorder : bool
+                Whether to reorder the data to be closest to canonical (RAS+) orientation.
         """
         cls.resample = resample
         cls.reorder = reorder
 
-    def show(self, ctx=None, channel=0, indices=None, anatomical_plane=0, **kwargs):
+    def show(self, ctx=None, channel: int = 0, indices: int = None, anatomical_plane: int = 0, **kwargs):
         """
-        Show Medimage using `merge(self._show_args, kwargs)`.
+        Displays the Medimage using `merge(self._show_args, kwargs)`.
+
+        Args:
+            ctx : Any, optional
+                Context to use for the display. Defaults to None.
+            channel : int, optional
+                The channel of the image to be displayed. Defaults to 0.
+            indices : list or None, optional
+                Indices of the images to be displayed. Defaults to None.
+            anatomical_plane : int, optional
+                Anatomical plane of the image to be displayed. Defaults to 0.
+            kwargs : dict, optional
+                Additional parameters for the show function.
 
         Returns:
             Shown image.
@@ -177,15 +192,16 @@ class MedBase(torch.Tensor, metaclass=MetaResolver):
             **merge(self._show_args, kwargs)
         )
 
-    def __repr__(self):
+    def __repr__(self) -> str:
+        """Returns the string representation of the MedBase instance."""
         return f'{self.__class__.__name__} mode={self.mode} size={"x".join([str(d) for d in self.size])}'
 
 # %% ../nbs/01_vision_core.ipynb 12
 class MedImage(MedBase):
-    '''Subclass of MedBase that represents an image object.'''
+    """Subclass of MedBase that represents an image object."""
     pass
 
 # %% ../nbs/01_vision_core.ipynb 13
 class MedMask(MedBase):
-    '''Subclass of MedBase that represents an mask object.'''
+    """Subclass of MedBase that represents an mask object."""
     _show_args = {'alpha':0.5, 'cmap':'tab20'}