careamics 0.1.0rc5__py3-none-any.whl → 0.1.0rc6__py3-none-any.whl

careamics/dataset/patching/sequential_patching.py

@@ -1,3 +1,5 @@
+"""Sequential patching functions."""
+
 from typing import List, Optional, Tuple, Union
 
 import numpy as np
@@ -14,14 +16,14 @@ def _compute_number_of_patches(
 
     Parameters
     ----------
-    arr : Tuple[int, ...]
+    arr_shape : Tuple[int, ...]
         Shape of the input array.
-    patch_sizes : Tuple[int]
+    patch_sizes : Union[List[int], Tuple[int, ...]
         Shape of the patches.
 
     Returns
     -------
-    Tuple[int]
+    Tuple[int, ...]
         Number of patches in each dimension.
     """
     if len(arr_shape) != len(patch_sizes):
@@ -55,14 +57,14 @@ def _compute_overlap(
 
     Parameters
     ----------
-    arr : Tuple[int, ...]
+    arr_shape : Tuple[int, ...]
         Input array shape.
-    patch_sizes : Tuple[int]
+    patch_sizes : Union[List[int], Tuple[int, ...]]
         Size of the patches.
 
     Returns
     -------
-    Tuple[int]
+    Tuple[int, ...]
         Overlap between patches in each dimension.
     """
     n_patches = _compute_number_of_patches(arr_shape, patch_sizes)
@@ -123,6 +125,8 @@ def _compute_patch_views(
         Steps between views.
     output_shape : Tuple[int]
         Shape of the output array.
+    target : Optional[np.ndarray], optional
+        Target array, by default None.
 
     Returns
     -------
@@ -161,11 +165,13 @@ def extract_patches_sequential(
         Input image array.
     patch_size : Tuple[int]
         Patch sizes in each dimension.
+    target : Optional[np.ndarray], optional
+        Target array, by default None.
 
     Returns
     -------
-    Generator[Tuple[np.ndarray, ...], None, None]
-        Generator of patches.
+    Tuple[np.ndarray, Optional[np.ndarray]]
+        Patches.
     """
     is_3d_patch = len(patch_size) == 3
 

careamics/dataset/patching/tiled_patching.py

@@ -1,3 +1,5 @@
+"""Tiled patching utilities."""
+
 import itertools
 from typing import Generator, List, Tuple, Union
 
@@ -8,7 +10,7 @@ from careamics.config.tile_information import TileInformation
 
 def _compute_crop_and_stitch_coords_1d(
     axis_size: int, tile_size: int, overlap: int
-) -> Tuple[List[Tuple[int, ...]], ...]:
+) -> Tuple[List[Tuple[int, int]], List[Tuple[int, int]], List[Tuple[int, int]]]:
     """
     Compute the coordinates of each tile along an axis, given the overlap.
 

careamics/dataset/patching/validate_patch_dimension.py

@@ -1,3 +1,5 @@
+"""Patch validation functions."""
+
 from typing import List, Tuple, Union
 
 import numpy as np

careamics/dataset/zarr_dataset.py

@@ -1,3 +1,5 @@
+"""Zarr dataset."""
+
 # from itertools import islice
 # from typing import Callable, Dict, List, Optional, Tuple, Union
 

careamics/lightning_datamodule.py

@@ -95,13 +95,13 @@ class CAREamicsTrainData(L.LightningDataModule):
         Batch size.
     use_in_memory : bool
         Whether to use in memory dataset if possible.
-    train_data : Union[Path, str, np.ndarray]
+    train_data : Union[Path, np.ndarray]
         Training data.
-    val_data : Optional[Union[Path, str, np.ndarray]]
+    val_data : Optional[Union[Path, np.ndarray]]
         Validation data.
-    train_data_target : Optional[Union[Path, str, np.ndarray]]
+    train_data_target : Optional[Union[Path, np.ndarray]]
         Training target data.
-    val_data_target : Optional[Union[Path, str, np.ndarray]]
+    val_data_target : Optional[Union[Path, np.ndarray]]
         Validation target data.
     val_percentage : float
         Percentage of the training data to use for validation, if no validation data is
@@ -217,17 +217,33 @@ class CAREamicsTrainData(L.LightningDataModule):
             )
 
         # configuration
-        self.data_config = data_config
-        self.data_type = data_config.data_type
-        self.batch_size = data_config.batch_size
-        self.use_in_memory = use_in_memory
+        self.data_config: DataConfig = data_config
+        self.data_type: str = data_config.data_type
+        self.batch_size: int = data_config.batch_size
+        self.use_in_memory: bool = use_in_memory
+
+        # data: make data Path or np.ndarray, use type annotations for mypy
+        self.train_data: Union[Path, np.ndarray] = (
+            Path(train_data) if isinstance(train_data, str) else train_data
+        )
+
+        self.val_data: Union[Path, np.ndarray] = (
+            Path(val_data) if isinstance(val_data, str) else val_data
+        )
+
+        self.train_data_target: Union[Path, np.ndarray] = (
+            Path(train_data_target)
+            if isinstance(train_data_target, str)
+            else train_data_target
+        )
 
-        # data
-        self.train_data = train_data
-        self.val_data = val_data
+        self.val_data_target: Union[Path, np.ndarray] = (
+            Path(val_data_target)
+            if isinstance(val_data_target, str)
+            else val_data_target
+        )
 
-        self.train_data_target = train_data_target
-        self.val_data_target = val_data_target
+        # validation split
         self.val_percentage = val_percentage
         self.val_minimum_split = val_minimum_split
 
@@ -241,10 +257,10 @@ class CAREamicsTrainData(L.LightningDataModule):
         elif data_config.data_type != SupportedData.ARRAY:
             self.read_source_func = get_read_func(data_config.data_type)
 
-        self.extension_filter = extension_filter
+        self.extension_filter: str = extension_filter
 
         # Pytorch dataloader parameters
-        self.dataloader_params = (
+        self.dataloader_params: Dict[str, Any] = (
             data_config.dataloader_params if data_config.dataloader_params else {}
         )
 
@@ -309,20 +325,30 @@ class CAREamicsTrainData(L.LightningDataModule):
         """
         # if numpy array
         if self.data_type == SupportedData.ARRAY:
+            # mypy checks
+            assert isinstance(self.train_data, np.ndarray)
+            if self.train_data_target is not None:
+                assert isinstance(self.train_data_target, np.ndarray)
+
             # train dataset
             self.train_dataset: DatasetType = InMemoryDataset(
                 data_config=self.data_config,
                 inputs=self.train_data,
-                data_target=self.train_data_target,
+                input_target=self.train_data_target,
             )
 
             # validation dataset
             if self.val_data is not None:
+                # mypy checks
+                assert isinstance(self.val_data, np.ndarray)
+                if self.val_data_target is not None:
+                    assert isinstance(self.val_data_target, np.ndarray)
+
                 # create its own dataset
                 self.val_dataset: DatasetType = InMemoryDataset(
                     data_config=self.data_config,
                     inputs=self.val_data,
-                    data_target=self.val_data_target,
+                    input_target=self.val_data_target,
                 )
             else:
                 # extract validation from the training patches
@@ -341,7 +367,7 @@ class CAREamicsTrainData(L.LightningDataModule):
                 self.train_dataset = InMemoryDataset(
                     data_config=self.data_config,
                     inputs=self.train_files,
-                    data_target=(
+                    input_target=(
                         self.train_target_files if self.train_data_target else None
                     ),
                     read_source_func=self.read_source_func,
@@ -352,7 +378,7 @@ class CAREamicsTrainData(L.LightningDataModule):
                     self.val_dataset = InMemoryDataset(
                         data_config=self.data_config,
                         inputs=self.val_files,
-                        data_target=(
+                        input_target=(
                             self.val_target_files if self.val_data_target else None
                         ),
                         read_source_func=self.read_source_func,

careamics/lightning_module.py

@@ -1,3 +1,5 @@
+"""CAREamics Lightning module."""
+
 from typing import Any, Optional, Union
 
 import pytorch_lightning as L
@@ -24,6 +26,11 @@ class CAREamicsModule(L.LightningModule):
     This class encapsulates the a PyTorch model along with the training, validation,
     and testing logic. It is configured using an `AlgorithmModel` Pydantic class.
 
+    Parameters
+    ----------
+    algorithm_config : Union[AlgorithmModel, dict]
+        Algorithm configuration.
+
     Attributes
     ----------
     model : nn.Module
@@ -39,8 +46,7 @@ class CAREamicsModule(L.LightningModule):
     """
 
     def __init__(self, algorithm_config: Union[AlgorithmConfig, dict]) -> None:
-        """
-        CAREamics Lightning module.
+        """Lightning module for CAREamics.
 
         This class encapsulates the a PyTorch model along with the training, validation,
         and testing logic. It is configured using an `AlgorithmModel` Pydantic class.

careamics/lightning_prediction_datamodule.py

@@ -1,7 +1,7 @@
 """Prediction Lightning data modules."""
 
 from pathlib import Path
-from typing import Any, Callable, List, Literal, Optional, Tuple, Union
+from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
 
 import numpy as np
 import pytorch_lightning as L
@@ -303,9 +303,6 @@ class PredictDataWrapper(CAREamicsPredictData):
         Batch size.
     tta_transforms : bool, optional
         Use test time augmentation, by default True.
-    transforms : List, optional
-        List of transforms to apply to prediction patches. If None, default
-        transforms are applied.
     read_source_func : Optional[Callable], optional
         Function to read the source data, used if `data_type` is `custom`, by
         default None.
@@ -326,7 +323,6 @@ class PredictDataWrapper(CAREamicsPredictData):
         axes: str = "YX",
         batch_size: int = 1,
         tta_transforms: bool = True,
-        transforms: Optional[List] = None,
         read_source_func: Optional[Callable] = None,
         extension_filter: str = "",
         dataloader_params: Optional[dict] = None,
@@ -356,9 +352,6 @@ class PredictDataWrapper(CAREamicsPredictData):
             Batch size.
         tta_transforms : bool, optional
             Use test time augmentation, by default True.
-        transforms : Optional[List], optional
-            List of transforms to apply to prediction patches. If None, default
-            transforms are applied.
         read_source_func : Optional[Callable], optional
             Function to read the source data, used if `data_type` is `custom`, by
             default None.
@@ -369,7 +362,7 @@ class PredictDataWrapper(CAREamicsPredictData):
         """
         if dataloader_params is None:
             dataloader_params = {}
-        prediction_dict = {
+        prediction_dict: Dict[str, Any] = {
             "data_type": data_type,
             "tile_size": tile_size,
             "tile_overlap": tile_overlap,
@@ -378,12 +371,9 @@ class PredictDataWrapper(CAREamicsPredictData):
             "std": std,
             "tta": tta_transforms,
             "batch_size": batch_size,
+            "transforms": [],
         }
 
-        # if transforms are passed (otherwise it will use the default ones)
-        if transforms is not None:
-            prediction_dict["transforms"] = transforms
-
         # validate configuration
         self.prediction_config = InferenceConfig(**prediction_dict)
 

careamics/lightning_prediction_loop.py

@@ -1,3 +1,5 @@
+"""Lithning prediction loop allowing tiling."""
+
 from typing import Optional
 
 import pytorch_lightning as L
@@ -18,14 +20,14 @@ class CAREamicsPredictionLoop(L.loops._PredictionLoop):
     """
 
     def _on_predict_epoch_end(self) -> Optional[_PREDICT_OUTPUT]:
-        """
-        Calls `on_predict_epoch_end` hook.
+        """Call `on_predict_epoch_end` hook.
 
         Adapted from the parent method.
 
         Returns
         -------
-            the results for all dataloaders
+        Optional[_PREDICT_OUTPUT]
+            Prediction output.
         """
         trainer = self.trainer
         call._call_callback_hooks(trainer, "on_predict_epoch_end")
@@ -45,15 +47,14 @@ class CAREamicsPredictionLoop(L.loops._PredictionLoop):
 
     @_no_grad_context
     def run(self) -> Optional[_PREDICT_OUTPUT]:
-        """
-        Runs the prediction loop.
+        """Run the prediction loop.
 
         Adapted from the parent method in order to stitch the predictions.
 
         Returns
         -------
         Optional[_PREDICT_OUTPUT]
-            Prediction output
+            Prediction output.
         """
         self.setup_data()
         if self.skip:
@@ -86,6 +87,7 @@ class CAREamicsPredictionLoop(L.loops._PredictionLoop):
 
                 ########################################################
                 ################ CAREamics specific code ###############
+                # TODO: next line is not compatible with muSplit
                 is_tiled = len(self.predictions[batch_idx]) == 2
                 if is_tiled:
                     # extract the last tile flag and the coordinates (crop and stitch)

careamics/losses/__init__.py

@@ -1,6 +1,5 @@
 """Losses module."""
 
-from .loss_factory import loss_factory
+__all__ = ["loss_factory"]
 
-# from .noise_model_factory import noise_model_factory as noise_model_factory
-# from .noise_models import GaussianMixtureNoiseModel, HistogramNoiseModel
+from .loss_factory import loss_factory

careamics/losses/loss_factory.py

@@ -17,7 +17,7 @@ def loss_factory(loss: Union[SupportedLoss, str]) -> Callable:
 
     Parameters
     ----------
-    loss: SupportedLoss
+    loss : Union[SupportedLoss, str]
         Requested loss.
 
     Returns

careamics/losses/losses.py

@@ -5,23 +5,27 @@ This submodule contains the various losses used in CAREamics.
 """
 
 import torch
-
-# TODO if we are only using the DiceLoss, can we just implement it?
-# from segmentation_models_pytorch.losses import DiceLoss
 from torch.nn import L1Loss, MSELoss
 
 
-def mse_loss(samples: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
+def mse_loss(source: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
     """
     Mean squared error loss.
 
+    Parameters
+    ----------
+    source : torch.Tensor
+        Source patches.
+    target : torch.Tensor
+        Target patches.
+
     Returns
     -------
     torch.Tensor
         Loss value.
     """
     loss = MSELoss()
-    return loss(samples, labels)
+    return loss(source, target)
 
 
 def n2v_loss(
@@ -34,9 +38,9 @@ def n2v_loss(
 
     Parameters
     ----------
-    samples : torch.Tensor
+    manipulated_patches : torch.Tensor
         Patches with manipulated pixels.
-    labels : torch.Tensor
+    original_patches : torch.Tensor
         Noisy patches.
     masks : torch.Tensor
         Array containing masked pixel locations.

careamics/model_io/bmz_io.py

@@ -104,9 +104,9 @@ def export_to_bmz(
     authors : List[dict]
         Authors of the model.
     input_array : np.ndarray
-        Input array.
+        Input array, should not have been normalized.
     output_array : np.ndarray
-        Output array.
+        Output array, should have been denormalized.
     channel_names : Optional[List[str]], optional
         Channel names, by default None.
     data_description : Optional[str], optional
@@ -178,7 +178,7 @@ def export_to_bmz(
         )
 
         # test model description
-        summary: ValidationSummary = test_model(model_description, decimal=0)
+        summary: ValidationSummary = test_model(model_description, decimal=2)
         if summary.status == "failed":
             raise ValueError(f"Model description test failed: {summary}")
 

careamics/models/activation.py

@@ -1,3 +1,5 @@
+"""Activations for CAREamics models."""
+
 from typing import Callable, Union
 
 import torch.nn as nn

careamics/models/layers.py

@@ -162,6 +162,18 @@ def _unpack_kernel_size(
     """Unpack kernel_size to a tuple of ints.
 
     Inspired by Kornia implementation. TODO: link
+
+    Parameters
+    ----------
+    kernel_size : Union[Tuple[int, ...], int]
+        Kernel size.
+    dim : int
+        Number of dimensions.
+
+    Returns
+    -------
+    Tuple[int, ...]
+        Kernel size tuple.
     """
     if isinstance(kernel_size, int):
         kernel_dims = tuple([kernel_size for _ in range(dim)])
@@ -173,7 +185,20 @@ def _unpack_kernel_size(
 def _compute_zero_padding(
     kernel_size: Union[Tuple[int, ...], int], dim: int
 ) -> Tuple[int, ...]:
-    """Utility function that computes zero padding tuple."""
+    """Utility function that computes zero padding tuple.
+
+    Parameters
+    ----------
+    kernel_size : Union[Tuple[int, ...], int]
+        Kernel size.
+    dim : int
+        Number of dimensions.
+
+    Returns
+    -------
+    Tuple[int, ...]
+        Zero padding tuple.
+    """
     kernel_dims = _unpack_kernel_size(kernel_size, dim)
     return tuple([(kd - 1) // 2 for kd in kernel_dims])
 
@@ -191,14 +216,19 @@ def get_pascal_kernel_1d(
 
     Parameters
     ----------
-    kernel_size: height and width of the kernel.
-    norm: if to normalize the kernel or not. Default: False.
-    device: tensor device
-    dtype: tensor dtype
+    kernel_size : int
+        Kernel size.
+    norm : bool
+        Normalize the kernel, by default False.
+    device : Optional[torch.device]
+        Device of the tensor, by default None.
+    dtype : Optional[torch.dtype]
+        Data type of the tensor, by default None.
 
     Returns
     -------
-    kernel shaped as :math:`(kernel_size,)`
+    torch.Tensor
+        Pascal kernel.
 
     Examples
     --------
@@ -245,19 +275,28 @@ def _get_pascal_kernel_nd(
 ) -> torch.Tensor:
     """Generate pascal filter kernel by kernel size.
 
+    If kernel_size is an integer the kernel will be shaped as (kernel_size, kernel_size)
+    otherwise the kernel will be shaped as kernel_size
+
     Inspired by Kornia implementation.
 
     Parameters
     ----------
-    kernel_size: height and width of the kernel.
-    norm: if to normalize the kernel or not. Default: True.
-    device: tensor device
-    dtype: tensor dtype
+    kernel_size : Union[Tuple[int, int], int]
+        Kernel size for the pascal kernel.
+    norm : bool
+        Normalize the kernel, by default True.
+    dim : int
+        Number of dimensions, by default 2.
+    device : Optional[torch.device]
+        Device of the tensor, by default None.
+    dtype : Optional[torch.dtype]
+        Data type of the tensor, by default None.
 
     Returns
     -------
-    if kernel_size is an integer the kernel will be shaped as (kernel_size, kernel_size)
-    otherwise the kernel will be shaped as kernel_size
+    torch.Tensor
+        Pascal kernel.
 
     Examples
     --------
@@ -303,6 +342,24 @@ def _max_blur_pool_by_kernel2d(
     """Compute max_blur_pool by a given :math:`CxC_(out, None)xNxN` kernel.
 
     Inspired by Kornia implementation.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Input tensor.
+    kernel : torch.Tensor
+        Kernel tensor.
+    stride : int
+        Stride.
+    max_pool_size : int
+        Maximum pool size.
+    ceil_mode : bool
+        Ceil mode, by default False. Set to True to match output size of conv2d.
+
+    Returns
+    -------
+    torch.Tensor
+        Output tensor.
     """
     # compute local maxima
     x = F.max_pool2d(
@@ -323,6 +380,24 @@ def _max_blur_pool_by_kernel3d(
     """Compute max_blur_pool by a given :math:`CxC_(out, None)xNxNxN` kernel.
 
     Inspired by Kornia implementation.
+
+    Parameters
+    ----------
+    x : torch.Tensor
+        Input tensor.
+    kernel : torch.Tensor
+        Kernel tensor.
+    stride : int
+        Stride.
+    max_pool_size : int
+        Maximum pool size.
+    ceil_mode : bool
+        Ceil mode, by default False. Set to True to match output size of conv2d.
+
+    Returns
+    -------
+    torch.Tensor
+        Output tensor.
     """
     # compute local maxima
     x = F.max_pool3d(
@@ -343,21 +418,16 @@ class MaxBlurPool(nn.Module):
 
     Parameters
     ----------
-    dim: int
-        Toggles between 2D and 3D
-    kernel_size: Union[Tuple[int, int], int]
+    dim : int
+        Toggles between 2D and 3D.
+    kernel_size : Union[Tuple[int, int], int]
         Kernel size for max pooling.
-    stride: int
+    stride : int
         Stride for pooling.
-    max_pool_size: int
+    max_pool_size : int
         Max kernel size for max pooling.
-    ceil_mode: bool
-        Should be true to match output size of conv2d with same kernel size.
-
-    Returns
-    -------
-    torch.Tensor
-        The pooled and blurred tensor.
+    ceil_mode : bool
+        Ceil mode, by default False. Set to True to match output size of conv2d.
     """
 
     def __init__(
@@ -368,6 +438,21 @@ class MaxBlurPool(nn.Module):
         max_pool_size: int = 2,
         ceil_mode: bool = False,
     ) -> None:
+        """Constructor.
+
+        Parameters
+        ----------
+        dim : int
+            Dimension of the convolution.
+        kernel_size : Union[Tuple[int, int], int]
+            Kernel size for max pooling.
+        stride : int, optional
+            Stride, by default 2.
+        max_pool_size : int, optional
+            Maximum pool size, by default 2.
+        ceil_mode : bool, optional
+            Ceil mode, by default False. Set to True to match output size of conv2d.
+        """
         super().__init__()
         self.dim = dim
         self.kernel_size = kernel_size
@@ -377,7 +462,18 @@ class MaxBlurPool(nn.Module):
         self.kernel = _get_pascal_kernel_nd(kernel_size, norm=True, dim=self.dim)
 
     def forward(self, x: torch.Tensor) -> torch.Tensor:
-        """Forward pass of the function."""
+        """Forward pass of the function.
+
+        Parameters
+        ----------
+        x : torch.Tensor
+            Input tensor.
+
+        Returns
+        -------
+        torch.Tensor
+            Output tensor.
+        """
         self.kernel = torch.as_tensor(self.kernel, device=x.device, dtype=x.dtype)
         if self.dim == 2:
             return _max_blur_pool_by_kernel2d(

careamics/models/model_factory.py

@@ -27,7 +27,7 @@ def model_factory(
     Parameters
     ----------
     model_configuration : Union[UNetModel, VAEModel]
-        Model configuration
+        Model configuration.
 
     Returns
     -------