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

fastMONAI/vision_data.py

@@ -10,168 +10,258 @@ from fastai.vision.data import *
 from .vision_core import *
 
 # %% ../nbs/02_vision_data.ipynb 5
-def pred_to_multiclass_mask(pred:torch.Tensor # [C,W,H,D] activation tensor
-                           ) -> torch.Tensor:
-    '''Apply Softmax function on the predicted tensor to rescale the values in the range [0, 1] and sum to 1.
-    Then apply argmax to get the indices of the maximum value of all elements in the predicted Tensor.
-    Returns: Predicted mask.
-    '''
+def pred_to_multiclass_mask(pred: torch.Tensor) -> torch.Tensor:
+    """Apply Softmax on the predicted tensor to rescale the values in the range [0, 1]
+    and sum to 1. Then apply argmax to get the indices of the maximum value of all 
+    elements in the predicted Tensor.
+
+    Args:
+        pred: [C,W,H,D] activation tensor.
+
+    Returns: 
+        Predicted mask.
+    """
+    
     pred = pred.softmax(dim=0)
+
     return pred.argmax(dim=0, keepdims=True)
 
 # %% ../nbs/02_vision_data.ipynb 6
-def batch_pred_to_multiclass_mask(pred:torch.Tensor # [B, C, W, H, D] batch of activations
-                                 ) -> (torch.Tensor, int):
-    '''Convert a batch of predicted activation tensors to masks.
-    Returns batch of predicted masks and number of classes.
-    '''
-
+def batch_pred_to_multiclass_mask(pred: torch.Tensor) -> (torch.Tensor, int):
+    """Convert a batch of predicted activation tensors to masks.
+    
+    Args:
+        pred: [B, C, W, H, D] batch of activations.
+
+    Returns:
+        Tuple of batch of predicted masks and number of classes.
+    """
+    
     n_classes = pred.shape[1]
     pred = [pred_to_multiclass_mask(p) for p in pred]
 
     return torch.stack(pred), n_classes
 
 # %% ../nbs/02_vision_data.ipynb 7
-def pred_to_binary_mask(pred # [B, C, W, H, D] or [C, W, H, D] activation tensor
-                       ) -> torch.Tensor:
-    '''Apply Sigmoid function that squishes activations into a range between 0 and 1.
-    Then we classify all values greater than or equal to 0.5 to 1, and the values below 0.5 to 0.
-
-    Returns predicted binary mask(s).
-    '''
-
+def pred_to_binary_mask(pred: torch.Tensor) -> torch.Tensor:
+    """Apply Sigmoid function that squishes activations into a range between 0 and 1.
+    Then we classify all values greater than or equal to 0.5 to 1, 
+    and the values below 0.5 to 0.
+
+    Args:
+        pred: [B, C, W, H, D] or [C, W, H, D] activation tensor
+
+    Returns:
+        Predicted binary mask(s).
+    """
+    
     pred = torch.sigmoid(pred)
-    return torch.where(pred>=0.5, 1, 0)
+
+    return torch.where(pred >= 0.5, 1, 0)
 
 # %% ../nbs/02_vision_data.ipynb 9
 class MedDataBlock(DataBlock):
-    '''Container to quickly build dataloaders.'''
+    """Container to quickly build dataloaders."""
+    #TODO add get_x
+    def __init__(self, blocks: list = None, dl_type: TfmdDL = None, getters: list = None,
+                 n_inp: int = None, item_tfms: list = None, batch_tfms: list = None,
+                 reorder: bool = False, resample: (int, list) = None, **kwargs):
 
-    def __init__(self, blocks:list=None,dl_type:TfmdDL=None, getters:list=None, n_inp:int=None, item_tfms:list=None,
-                 batch_tfms:list=None, reorder:bool=False, resample:(int, list)=None, **kwargs):
+        super().__init__(blocks, dl_type, getters, n_inp, item_tfms,
+                         batch_tfms, **kwargs)
 
-        super().__init__(blocks, dl_type, getters, n_inp, item_tfms, batch_tfms, **kwargs)
-        MedBase.item_preprocessing(resample,reorder)
+        MedBase.item_preprocessing(resample, reorder)
 
-# %% ../nbs/02_vision_data.ipynb 12
+# %% ../nbs/02_vision_data.ipynb 11
 def MedMaskBlock():
+    """Create a TransformBlock for medical masks."""
     return TransformBlock(type_tfms=MedMask.create)
 
-# %% ../nbs/02_vision_data.ipynb 14
+# %% ../nbs/02_vision_data.ipynb 13
 class MedImageDataLoaders(DataLoaders):
-    '''Higher-level `MedDataBlock` API.'''
-
+    """Higher-level `MedDataBlock` API."""
+    
     @classmethod
     @delegates(DataLoaders.from_dblock)
-    def from_df(cls, df, valid_pct=0.2, seed=None, fn_col=0, folder=None, suff='', label_col=1, label_delim=None,
-                y_block=None, valid_col=None, item_tfms=None, batch_tfms=None, reorder=False, resample=None, **kwargs):
-        '''Create from DataFrame.'''
-
+    def from_df(cls, df, valid_pct=0.2, seed=None, fn_col=0, folder=None, suff='',
+                label_col=1, label_delim=None, y_block=None, valid_col=None,
+                item_tfms=None, batch_tfms=None, reorder=False, resample=None, **kwargs):
+        """Create from DataFrame."""
+                    
         if y_block is None:
             is_multi = (is_listy(label_col) and len(label_col) > 1) or label_delim is not None
             y_block = MultiCategoryBlock if is_multi else CategoryBlock
-        splitter = RandomSplitter(valid_pct, seed=seed) if valid_col is None else ColSplitter(valid_col)
 
+        splitter = (RandomSplitter(valid_pct, seed=seed) 
+                    if valid_col is None else ColSplitter(valid_col))
 
-        dblock = MedDataBlock(blocks=(ImageBlock(cls=MedImage), y_block), get_x=ColReader(fn_col, suff=suff),
-                              get_y=ColReader(label_col, label_delim=label_delim),
-                              splitter=splitter,
-                              item_tfms=item_tfms,
-                              reorder=reorder,
-                              resample=resample)
+        dblock = MedDataBlock(
+            blocks=(ImageBlock(cls=MedImage), y_block),
+            get_x=ColReader(fn_col, suff=suff),
+            get_y=ColReader(label_col, label_delim=label_delim),
+            splitter=splitter,
+            item_tfms=item_tfms,
+            reorder=reorder,
+            resample=resample
+        )
 
         return cls.from_dblock(dblock, df, **kwargs)
 
-# %% ../nbs/02_vision_data.ipynb 19
+# %% ../nbs/02_vision_data.ipynb 16
 @typedispatch
-def show_batch(x:MedImage, y, samples, ctxs=None, max_n=6, nrows=None, ncols=None, figsize=None, channel=0, indices=None, anatomical_plane=0, **kwargs):
+def show_batch(x: MedImage, y, samples, ctxs=None, max_n=6, nrows=None, 
+               ncols=None, figsize=None, channel=0, indices=None,
+               anatomical_plane=0, **kwargs):
     '''Showing a batch of samples for classification and regression tasks.'''
 
-    if ctxs is None: ctxs = get_grid(min(len(samples), max_n), nrows=nrows, ncols=ncols, figsize=figsize)
+    if ctxs is None: 
+        ctxs = get_grid(min(len(samples), max_n), nrows=nrows, ncols=ncols, figsize=figsize)
+        
     n = 1 if y is None else 2
+    
     for i in range(n):
-        ctxs = [b.show(ctx=c, channel=channel, indices=indices, anatomical_plane=anatomical_plane, **kwargs) for b,c,_ in zip(samples.itemgot(i),ctxs,range(max_n))]
+        ctxs = [
+            b.show(ctx=c, channel=channel, indices=indices, anatomical_plane=anatomical_plane, **kwargs)
+            for b,c,_ in zip(samples.itemgot(i),ctxs,range(max_n))
+        ]
 
     plt.tight_layout()
+    
     return ctxs
 
-# %% ../nbs/02_vision_data.ipynb 20
+# %% ../nbs/02_vision_data.ipynb 17
 @typedispatch
-def show_batch(x:MedImage, y:MedMask, samples, ctxs=None, max_n=6, nrows=None, ncols=None, figsize=None, channel=0, indices=None, anatomical_plane=0, **kwargs):
-    '''Showing a batch of decoded segmentation samples.'''
+def show_batch(x: MedImage, y: MedMask, samples, ctxs=None, max_n=6, nrows=None,
+               ncols=None, figsize=None, channel=0, indices=None,
+               anatomical_plane=0, **kwargs):
+    """Showing a batch of decoded segmentation samples."""
 
     nrows, ncols = min(len(samples), max_n), x.shape[1] + 1
     imgs = []
 
-    fig,axs = subplots(nrows, ncols, figsize=figsize, **kwargs)
+    fig, axs = subplots(nrows, ncols, figsize=figsize, **kwargs)
     axs = axs.flatten()
 
-    for img, mask in list(zip(x,y)):
+    for img, mask in zip(x, y):
         im_channels = [MedImage(c_img[None]) for c_img in img]
         im_channels.append(MedMask(mask))
         imgs.extend(im_channels)
 
-    ctxs = [im.show(ax=ax, indices=indices, anatomical_plane=anatomical_plane) for im, ax in zip(imgs, axs)]
+    ctxs = [im.show(ax=ax, indices=indices, anatomical_plane=anatomical_plane)
+            for im, ax in zip(imgs, axs)]
+
     plt.tight_layout()
 
     return ctxs
 
-# %% ../nbs/02_vision_data.ipynb 22
+# %% ../nbs/02_vision_data.ipynb 19
 @typedispatch
-def show_results(x:MedImage, y:torch.Tensor, samples, outs, ctxs=None, max_n=6, nrows=None, ncols=None, figsize=None, channel=0, indices=None, anatomical_plane=0, **kwargs):
-    '''Showing samples and their corresponding predictions for regression tasks.'''
+def show_results(x: MedImage, y: torch.Tensor, samples, outs, ctxs=None, max_n: int = 6,
+                 nrows: int = None, ncols: int = None, figsize=None, channel: int = 0,
+                 indices: int = None, anatomical_plane: int = 0, **kwargs):
+    """Showing samples and their corresponding predictions for regression tasks."""
 
-    if ctxs is None: ctxs = get_grid(min(len(samples), max_n), nrows=nrows, ncols=ncols, figsize=figsize)
+    if ctxs is None:
+        ctxs = get_grid(min(len(samples), max_n), nrows=nrows,
+                        ncols=ncols, figsize=figsize)
 
     for i in range(len(samples[0])):
-        ctxs = [b.show(ctx=c, channel=channel, indices=indices, anatomical_plane=anatomical_plane, **kwargs) for b,c,_ in zip(samples.itemgot(i),ctxs,range(max_n))]
+        ctxs = [
+            b.show(ctx=c, channel=channel, indices=indices,
+                   anatomical_plane=anatomical_plane, **kwargs)
+            for b, c, _ in zip(samples.itemgot(i), ctxs, range(max_n))
+        ]
+
     for i in range(len(outs[0])):
-        ctxs = [b.show(ctx=c, **kwargs) for b,c,_ in zip(outs.itemgot(i),ctxs,range(max_n))]
+        ctxs = [
+            b.show(ctx=c, **kwargs)
+            for b, c, _ in zip(outs.itemgot(i), ctxs, range(max_n))
+        ]
+
     return ctxs
 
-# %% ../nbs/02_vision_data.ipynb 23
+# %% ../nbs/02_vision_data.ipynb 20
 @typedispatch
-def show_results(x:MedImage, y:TensorCategory, samples, outs, ctxs=None, max_n=6, nrows=None, ncols=None, figsize=None, channel=0, indices=None, anatomical_plane=0, **kwargs):
-    '''Showing samples and their corresponding predictions for classification tasks.'''
-
-    if ctxs is None: ctxs = get_grid(min(len(samples), max_n), nrows=nrows, ncols=ncols, figsize=figsize)
+def show_results(x: MedImage, y: TensorCategory, samples, outs, ctxs=None, 
+                 max_n: int = 6, nrows: int = None, ncols: int = None, figsize=None, channel: int = 0, 
+                 indices: int = None, anatomical_plane: int = 0, **kwargs):
+    """Showing samples and their corresponding predictions for classification tasks."""
+
+    if ctxs is None: 
+        ctxs = get_grid(min(len(samples), max_n), nrows=nrows, 
+                        ncols=ncols, figsize=figsize)
+    
     for i in range(2):
-        ctxs = [b.show(ctx=c, channel=channel, indices=indices, anatomical_plane=anatomical_plane, **kwargs) for b,c,_ in zip(samples.itemgot(i),ctxs,range(max_n))]
-    ctxs = [r.show(ctx=c, color='green' if b==r else 'red', **kwargs) for b,r,c,_ in zip(samples.itemgot(1),outs.itemgot(0),ctxs,range(max_n))]
+        ctxs = [b.show(ctx=c, channel=channel, indices=indices, 
+                       anatomical_plane=anatomical_plane, **kwargs) 
+                for b, c, _ in zip(samples.itemgot(i), ctxs, range(max_n))]
+
+    ctxs = [r.show(ctx=c, color='green' if b == r else 'red', **kwargs) 
+            for b, r, c, _ in zip(samples.itemgot(1), outs.itemgot(0), ctxs, range(max_n))]
+
     return ctxs
 
-# %% ../nbs/02_vision_data.ipynb 24
+# %% ../nbs/02_vision_data.ipynb 21
 @typedispatch
-def show_results(x:MedImage, y:MedMask, samples, outs, ctxs=None, max_n=6, nrows=None, ncols=1, figsize=None, channel=0, indices=None, anatomical_plane=0, **kwargs):
-    ''' Showing decoded samples and their corresponding predictions for segmentation tasks.'''
+def show_results(x: MedImage, y: MedMask, samples, outs, ctxs=None, max_n: int = 6, 
+                 nrows: int = None, ncols: int = 1, figsize=None, channel: int = 0, 
+                 indices: int = None, anatomical_plane: int = 0, **kwargs):
+    """Showing decoded samples and their corresponding predictions for segmentation tasks."""
+
+    if ctxs is None: 
+        ctxs = get_grid(min(len(samples), max_n), nrows=nrows, ncols=ncols, 
+                        figsize=figsize, double=True, title='Target/Prediction')
 
-    if ctxs is None: ctxs = get_grid(min(len(samples), max_n), nrows=nrows, ncols=ncols, figsize=figsize, double=True, title='Target/Prediction')
     for i in range(2):
-        ctxs[::2] = [b.show(ctx=c, channel=channel, indices=indices, anatomical_plane=anatomical_plane, **kwargs) for b,c,_ in zip(samples.itemgot(i),ctxs[::2],range(2*max_n))]
-    for o in [samples,outs]:
-        ctxs[1::2] = [b.show(ctx=c, channel=channel, indices=indices, anatomical_plane=anatomical_plane, **kwargs) for b,c,_ in zip(o.itemgot(0),ctxs[1::2],range(2*max_n))]
+        ctxs[::2] = [b.show(ctx=c, channel=channel, indices=indices, 
+                            anatomical_plane=anatomical_plane, **kwargs) 
+                     for b, c, _ in zip(samples.itemgot(i), ctxs[::2], range(2 * max_n))]
+
+    for o in [samples, outs]:
+        ctxs[1::2] = [b.show(ctx=c, channel=channel, indices=indices, 
+                             anatomical_plane=anatomical_plane, **kwargs) 
+                      for b, c, _ in zip(o.itemgot(0), ctxs[1::2], range(2 * max_n))]
+
     return ctxs
 
-# %% ../nbs/02_vision_data.ipynb 26
+# %% ../nbs/02_vision_data.ipynb 23
 @typedispatch
-def plot_top_losses(x: MedImage, y, samples, outs, raws, losses, nrows=None, ncols=None, figsize=None, channel=0, indices=None, anatomical_plane=0, **kwargs):
-    '''Show images in top_losses along with their prediction, actual, loss, and probability of actual class.'''
+def plot_top_losses(x: MedImage, y: TensorCategory, samples, outs, raws, losses, nrows: int = None, 
+                    ncols: int = None, figsize=None, channel: int = 0, indices: int = None, 
+                    anatomical_plane: int = 0, **kwargs):
+    """Show images in top_losses along with their prediction, actual, loss, and probability of actual class."""
 
-    title = 'Prediction/Actual/Loss' if type(y) == torch.Tensor else 'Prediction/Actual/Loss/Probability'
+    title = 'Prediction/Actual/Loss' if isinstance(y, torch.Tensor) else 'Prediction/Actual/Loss/Probability'
     axs = get_grid(len(samples), nrows=nrows, ncols=ncols, figsize=figsize, title=title)
-    for ax,s,o,r,l in zip(axs, samples, outs, raws, losses):
+
+    for ax, s, o, r, l in zip(axs, samples, outs, raws, losses):
         s[0].show(ctx=ax, channel=channel, indices=indices, anatomical_plane=anatomical_plane, **kwargs)
-        if type(y) == torch.Tensor: ax.set_title(f'{r.max().item():.2f}/{s[1]} / {l.item():.2f}')
-        else: ax.set_title(f'{o[0]}/{s[1]} / {l.item():.2f} / {r.max().item():.2f}')
 
-# %% ../nbs/02_vision_data.ipynb 27
+        if isinstance(y, torch.Tensor): 
+            ax.set_title(f'{r.max().item():.2f}/{s[1]} / {l.item():.2f}')
+        else: 
+            ax.set_title(f'{o[0]}/{s[1]} / {l.item():.2f} / {r.max().item():.2f}')
+
+# %% ../nbs/02_vision_data.ipynb 24
 @typedispatch
-def plot_top_losses(x: MedImage, y:TensorMultiCategory, samples, outs, raws, losses, nrows=None, ncols=None, figsize=None, channel=0, indices=None, anatomical_plane=0, **kwargs):
-    #TODO: not tested yet
+def plot_top_losses(x: MedImage, y: TensorMultiCategory, samples, outs, raws, 
+                    losses, nrows: int = None, ncols: int = None, figsize=None, 
+                    channel: int = 0, indices: int = None, 
+                    anatomical_plane: int = 0, **kwargs):
+    # TODO: not tested yet
     axs = get_grid(len(samples), nrows=nrows, ncols=ncols, figsize=figsize)
-    for i,(ax,s) in enumerate(zip(axs, samples)): s[0].show(ctx=ax, title=f'Image {i}', channel=channel, indices=indices, anatomical_plane=anatomical_plane, **kwargs)
+
+    for i, (ax, s) in enumerate(zip(axs, samples)):
+        s[0].show(ctx=ax, title=f'Image {i}', channel=channel, 
+                  indices=indices, anatomical_plane=anatomical_plane, **kwargs)
+
     rows = get_empty_df(len(samples))
-    outs = L(s[1:] + o + (TitledStr(r), TitledFloat(l.item())) for s,o,r,l in zip(samples, outs, raws, losses))
-    for i,l in enumerate(["target", "predicted", "probabilities", "loss"]):
-        rows = [b.show(ctx=r, label=l, channel=channel, indices=indices, anatomical_plane=anatomical_plane, **kwargs) for b,r in zip(outs.itemgot(i),rows)]
+    outs = L(s[1:] + o + (TitledStr(r), TitledFloat(l.item())) 
+             for s, o, r, l in zip(samples, outs, raws, losses))
+
+    for i, l in enumerate(["target", "predicted", "probabilities", "loss"]):
+        rows = [b.show(ctx=r, label=l, channel=channel, indices=indices, 
+                       anatomical_plane=anatomical_plane, **kwargs) 
+                for b, r in zip(outs.itemgot(i), rows)]
+
     display_df(pd.DataFrame(rows))

fastMONAI/vision_inference.py

@@ -4,15 +4,16 @@
 __all__ = ['inference', 'refine_binary_pred_mask']
 
 # %% ../nbs/06_vision_inference.ipynb 1
-import numpy as np
+from copy import copy
 from pathlib import Path
-from torchio import Resize
+import torch
+import numpy as np
 from scipy.ndimage import label
-from .vision_core import *
-from .vision_augmentation import do_pad_or_crop
 from skimage.morphology import remove_small_objects
 from SimpleITK import DICOMOrient, GetArrayFromImage
-from copy import copy
+from torchio import Resize
+from .vision_core import *
+from .vision_augmentation import do_pad_or_crop
 
 # %% ../nbs/06_vision_inference.ipynb 3
 def _to_original_orientation(input_img, org_orientation):
@@ -24,27 +25,42 @@ def _to_original_orientation(input_img, org_orientation):
     return reoriented_array[None]
 
 # %% ../nbs/06_vision_inference.ipynb 4
-def _do_resize(o, target_shape, image_interpolation='linear', label_interpolation='nearest'):
-    '''Resample images so the output shape matches the given target shape.'''
+def _do_resize(o, target_shape, image_interpolation='linear', 
+               label_interpolation='nearest'):
+    """
+    Resample images so the output shape matches the given target shape.
+    """
 
-    resize = Resize(target_shape, image_interpolation=image_interpolation, label_interpolation=label_interpolation)
+    resize = Resize(
+        target_shape, 
+        image_interpolation=image_interpolation, 
+        label_interpolation=label_interpolation
+    )
+    
     return resize(o)
 
 # %% ../nbs/06_vision_inference.ipynb 5
-def inference(learn_inf, reorder, resample, fn:(Path,str)='', save_path:(str,Path)=None, org_img=None, input_img=None, org_size=None): 
-    '''Predict on new data using exported model'''         
+def inference(learn_inf, reorder, resample, fn: (str, Path) = '',
+              save_path: (str, Path) = None, org_img=None, input_img=None,
+              org_size=None): 
+    """Predict on new data using exported model."""         
+    
     if None in [org_img, input_img, org_size]: 
-        org_img, input_img, org_size = med_img_reader(fn, reorder, resample, only_tensor=False)
-    else: org_img, input_img = copy(org_img), copy(input_img)
+        org_img, input_img, org_size = med_img_reader(fn, reorder, resample, 
+                                                      only_tensor=False)
+    else: 
+        org_img, input_img = copy(org_img), copy(input_img)
     
-    pred, *_ = learn_inf.predict(input_img.data);
+    pred, *_ = learn_inf.predict(input_img.data)
     
-    pred_mask = do_pad_or_crop(pred.float(), input_img.shape[1:], padding_mode=0, mask_name=None)
+    pred_mask = do_pad_or_crop(pred.float(), input_img.shape[1:], padding_mode=0, 
+                               mask_name=None)
     input_img.set_data(pred_mask)
     
     input_img = _do_resize(input_img, org_size, image_interpolation='nearest')
     
-    reoriented_array = _to_original_orientation(input_img.as_sitk(), ('').join(org_img.orientation))
+    reoriented_array = _to_original_orientation(input_img.as_sitk(), 
+                                                ('').join(org_img.orientation))
     
     org_img.set_data(reoriented_array)
 
@@ -56,12 +72,10 @@ def inference(learn_inf, reorder, resample, fn:(Path,str)='', save_path:(str,Pat
     return org_img
 
 # %% ../nbs/06_vision_inference.ipynb 7
-def refine_binary_pred_mask(
-        pred_mask,
-        remove_size: (int, float) = None,
-        percentage: float = 0.2,
-        verbose: bool = False
-):
+def refine_binary_pred_mask(pred_mask, 
+                            remove_size: (int, float) = None,
+                            percentage: float = 0.2,
+                            verbose: bool = False) -> torch.Tensor:
     """Removes small objects from the predicted binary mask.
 
     Args:
@@ -74,6 +88,7 @@ def refine_binary_pred_mask(
     Returns:
         The processed mask with small objects removed.
     """
+                                
     labeled_mask, n_components = label(pred_mask)
 
     if verbose:
@@ -88,4 +103,4 @@ def refine_binary_pred_mask(
     processed_mask = remove_small_objects(
         labeled_mask, min_size=small_objects_threshold)
 
-    return np.where(processed_mask > 0, 1., 0.)
+    return torch.Tensor(processed_mask > 0).float()                          

fastMONAI/vision_loss.py

@@ -12,40 +12,49 @@ from torch.nn.modules.loss import _Loss
 
 # %% ../nbs/04_vision_loss_functions.ipynb 3
 class CustomLoss:
-    '''Wrapper to get show_results to work.'''
+    """A custom loss wrapper class for loss functions to allow them to work with
+    the 'show_results' method in fastai. 
+    """
 
     def __init__(self, loss_func):
+        """Constructs CustomLoss object."""
+        
         self.loss_func = loss_func
 
     def __call__(self, pred, targ):
-        if isinstance(pred, MedBase): pred, targ = torch.Tensor(pred.cpu()), torch.Tensor(targ.cpu().float())
+        """Computes the loss for given predictions and targets."""
+        
+        if isinstance(pred, MedBase):
+            pred, targ = torch.Tensor(pred.cpu()), torch.Tensor(targ.cpu().float())
+            
         return self.loss_func(pred, targ)
 
     def activation(self, x):
         return x
     
-    def decodes(self, x):
-        '''Converts model output to target format.
-
+    def decodes(self, x) -> torch.Tensor:
+        """Converts model output to target format.
+        
         Args:
-            x: Activations for each class [B, C, W, H, D]
+            x: Activations for each class with dimensions [B, C, W, H, D].
 
         Returns:
-            torch.Tensor: Predicted mask.
-        '''
-
+            The predicted mask.
+        """
+        
         n_classes = x.shape[1]
-        if n_classes == 1: x = pred_to_binary_mask(x)
-        else: x,_ = batch_pred_to_multiclass_mask(x)
+        if n_classes == 1: 
+            x = pred_to_binary_mask(x)
+        else: 
+            x,_ = batch_pred_to_multiclass_mask(x)
 
         return x
 
 # %% ../nbs/04_vision_loss_functions.ipynb 4
 class TverskyFocalLoss(_Loss):
     """
-    Compute both Dice loss and Focal Loss, and return the weighted sum of these two losses.
-    The details of Dice loss is shown in ``monai.losses.DiceLoss``.
-    The details of Focal Loss is shown in ``monai.losses.FocalLoss``.
+    Compute Tversky loss with a focus parameter, gamma, applied.
+    The details of Tversky loss is shown in ``monai.losses.TverskyLoss``.
     """
 
     def __init__(
@@ -54,45 +63,45 @@ class TverskyFocalLoss(_Loss):
         to_onehot_y: bool = False,
         sigmoid: bool = False,
         softmax: bool = False,
-        reduction: str = "mean",
         gamma: float = 2,
-        #focal_weight: (float, int, torch.Tensor) = None,
-        #lambda_dice: float = 1.0,
-        #lambda_focal: float = 1.0,
-        alpha = 0.5, 
-        beta = 0.99
-    ) -> None:
-
+        alpha: float = 0.5, 
+        beta: float = 0.99):
+        """
+        Args:
+            include_background: if to calculate loss for the background class.
+            to_onehot_y: whether to convert `y` into one-hot format.
+            sigmoid: if True, apply a sigmoid function to the prediction.
+            softmax: if True, apply a softmax function to the prediction.
+            gamma: the focal parameter, it modulates the loss with regards to 
+                how far the prediction is from target.
+            alpha: the weight of false positive in Tversky loss calculation.
+            beta: the weight of false negative in Tversky loss calculation.
+        """
+            
         super().__init__()
-        self.tversky = TverskyLoss(to_onehot_y=to_onehot_y, include_background=include_background, sigmoid=sigmoid, softmax=softmax, alpha=alpha, beta=beta)
-        #self.focal = FocalLoss(to_onehot_y=to_onehot_y, include_background=include_background, gamma=gamma, weight=focal_weight, reduction=reduction)
-        
-        #if lambda_dice < 0.0: raise ValueError("lambda_dice should be no less than 0.0.")
-        #if lambda_focal < 0.0: raise ValueError("lambda_focal should be no less than 0.0.")
-        #self.lambda_dice = lambda_dice
-        #self.lambda_focal = lambda_focal
-        self.to_onehot_y = to_onehot_y
+        self.tversky = TverskyLoss(
+            to_onehot_y=to_onehot_y, 
+            include_background=include_background, 
+            sigmoid=sigmoid, 
+            softmax=softmax, 
+            alpha=alpha, 
+            beta=beta
+        )
         self.gamma = gamma
-        self.include_background = include_background
 
     def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
         """
         Args:
-            input: the shape should be BNH[WD]. The input should be the original logits
-                due to the restriction of ``monai.losses.FocalLoss``.
-            target: the shape should be BNH[WD] or B1H[WD].
+            input: the shape should be [B, C, W, H, D]. The input should be the original logits.
+            target: the shape should be[B, C, W, H, D].
+
         Raises:
             ValueError: When number of dimensions for input and target are different.
-            ValueError: When number of channels for target is neither 1 nor the same as input.
         """
         if len(input.shape) != len(target.shape):
-            raise ValueError("the number of dimensions for input and target should be the same.")
-
-        n_pred_ch = input.shape[1]
+            raise ValueError("The number of dimensions for input and target should be the same.")
 
         tversky_loss = self.tversky(input, target)
-        #focal_loss = self.focal(input, target)
-        total_loss: torch.Tensor = 1 - ((1 - tversky_loss)**self.gamma) #tversky_loss
-        #print(total_loss,total_loss.shape)
-        #tversky_loss +  focal_loss
+        total_loss: torch.Tensor = 1 - ((1 - tversky_loss)**self.gamma)
+
         return total_loss