revice documentation of dataset module

mu_map/dataset/default.py

+"""
+Module containing the default MuMapDataset definition.
+A dataset that maps a reconstruction to an attenuation map.
+"""
 from __future__ import annotations
 
 import os

mu_map/dataset/normalization.py

 """
 Module containing normalization methods either as functions
-or transformers.
+or transforms.
 """
 from typing import Any, Callable, Optional, Tuple
 
@@ -19,7 +19,7 @@ def norm_max(tensor: Tensor) -> Tensor:
     return (tensor - tensor.min()) / (tensor.max() - tensor.min())
 
 
-def norm_mean(tensor: Tensor):
+def norm_mean(tensor: Tensor) -> Transform:
     """
     Perform mean normalization on a tensor.
 
@@ -28,7 +28,7 @@ def norm_mean(tensor: Tensor):
     return tensor / tensor.mean()
 
 
-def norm_gaussian(tensor: Tensor):
+def norm_gaussian(tensor: Tensor) -> Transform:
     """
     Perform Gaussian normalization on a tensor.
 

mu_map/dataset/patches.py

+"""
+Module containing the MuMapPatchDataset which transforms
+the MuMapDataset in a way that patches of the reconstructions
+and attenuation maps are returned.
+"""
 import math
 import random
 
@@ -55,6 +60,18 @@ class MuMapPatchDataset(MuMapDataset):
         self.generate_patches()
 
     def copy(self, split_name: str):
+        """
+        Copy the dataset and modify the split.
+
+        Parameters
+        ----------
+        split_name: str
+            the name of the split the copy is about
+
+        Returns
+        -------
+        MuMapPatchDataset
+        """
         kwargs = self.kwargs.copy()
         kwargs["split_name"] = split_name
         return MuMapPatchDataset(
@@ -105,6 +122,13 @@ class MuMapPatchDataset(MuMapDataset):
                 self.patches.append((_id, z, y, x, padding))
 
     def __getitem__(self, index: int):
+        """
+        Get a reconstruction attenuation map patch pair by index.
+
+        Parameters
+        ----------
+        index: int
+        """
         _id, z, y, x, padding = self.patches[index]
         ps = self.patch_size
         ps_z = self.patch_size_z

mu_map/dataset/transform.py

+"""
+Module containing transforms for dataset.
+This can be pre-processing steps such as padding or cropping and augmentations
+applied during training time.
+"""
 import math
 from typing import List, Tuple
 

mu_map/dataset/util.py

+"""
+Dataset utility method.
+Currently, the only available method aligns images at
+their central slice.
+"""
 from typing import Tuple
 
 import numpy as np