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

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.
 
@@ -43,9 +45,11 @@ def _compute_crop_and_stitch_coords_1d(
             stitch_coords.append(
                 (
                     i + overlap // 2 if i > 0 else 0,
-                    i + tile_size - overlap // 2
-                    if crop_coords[-1][1] < axis_size
-                    else axis_size,
+                    (
+                        i + tile_size - overlap // 2
+                        if crop_coords[-1][1] < axis_size
+                        else axis_size
+                    ),
                 )
             )
 
@@ -53,9 +57,11 @@ def _compute_crop_and_stitch_coords_1d(
             overlap_crop_coords.append(
                 (
                     overlap // 2 if i > 0 else 0,
-                    tile_size - overlap // 2
-                    if crop_coords[-1][1] < axis_size
-                    else tile_size,
+                    (
+                        tile_size - overlap // 2
+                        if crop_coords[-1][1] < axis_size
+                        else tile_size
+                    ),
                 )
             )
 

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

@@ -1,10 +1,10 @@
 """Training and validation Lightning data modules."""
+
 from pathlib import Path
 from typing import Any, Callable, Dict, List, Literal, Optional, Union
 
 import numpy as np
 import pytorch_lightning as L
-from albumentations import Compose
 from torch.utils.data import DataLoader
 
 from careamics.config import DataConfig
@@ -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
+        )
 
-        # data
-        self.train_data = train_data
-        self.val_data = val_data
+        self.val_data: Union[Path, np.ndarray] = (
+            Path(val_data) if isinstance(val_data, str) else val_data
+        )
 
-        self.train_data_target = train_data_target
-        self.val_data_target = val_data_target
+        self.train_data_target: Union[Path, np.ndarray] = (
+            Path(train_data_target)
+            if isinstance(train_data_target, str)
+            else train_data_target
+        )
+
+        self.val_data_target: Union[Path, np.ndarray] = (
+            Path(val_data_target)
+            if isinstance(val_data_target, str)
+            else 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,9 +367,9 @@ class CAREamicsTrainData(L.LightningDataModule):
                 self.train_dataset = InMemoryDataset(
                     data_config=self.data_config,
                     inputs=self.train_files,
-                    data_target=self.train_target_files
-                    if self.train_data_target
-                    else None,
+                    input_target=(
+                        self.train_target_files if self.train_data_target else None
+                    ),
                     read_source_func=self.read_source_func,
                 )
 
@@ -352,9 +378,9 @@ class CAREamicsTrainData(L.LightningDataModule):
                     self.val_dataset = InMemoryDataset(
                         data_config=self.data_config,
                         inputs=self.val_files,
-                        data_target=self.val_target_files
-                        if self.val_data_target
-                        else None,
+                        input_target=(
+                            self.val_target_files if self.val_data_target else None
+                        ),
                         read_source_func=self.read_source_func,
                     )
                 else:
@@ -370,9 +396,9 @@ class CAREamicsTrainData(L.LightningDataModule):
                 self.train_dataset = PathIterableDataset(
                     data_config=self.data_config,
                     src_files=self.train_files,
-                    target_files=self.train_target_files
-                    if self.train_data_target
-                    else None,
+                    target_files=(
+                        self.train_target_files if self.train_data_target else None
+                    ),
                     read_source_func=self.read_source_func,
                 )
 
@@ -382,9 +408,9 @@ class CAREamicsTrainData(L.LightningDataModule):
                     self.val_dataset = PathIterableDataset(
                         data_config=self.data_config,
                         src_files=self.val_files,
-                        target_files=self.val_target_files
-                        if self.val_data_target
-                        else None,
+                        target_files=(
+                            self.val_target_files if self.val_data_target else None
+                        ),
                         read_source_func=self.read_source_func,
                     )
                 elif len(self.train_files) <= self.val_minimum_split:
@@ -452,8 +478,7 @@ class TrainingDataWrapper(CAREamicsTrainData):
     In particular, N2V requires a specific transformation (N2V manipulates), which is
     not compatible with supervised training. The default transformations applied to the
     training patches are defined in `careamics.config.data_model`. To use different
-    transformations, pass a list of transforms or an albumentation `Compose` as
-    `transforms` parameter. See examples for more details.
+    transformations, pass a list of transforms. See examples for more details.
 
     By default, CAREamics only supports types defined in
     `careamics.config.support.SupportedData`. To read custom data types, you can set
@@ -488,7 +513,7 @@ class TrainingDataWrapper(CAREamicsTrainData):
         Batch size.
     val_data : Optional[Union[str, Path]], optional
         Validation data, by default None.
-    transforms : Optional[Union[List[TRANSFORMS_UNION], Compose]], optional
+    transforms : List[TRANSFORMS_UNION], optional
         List of transforms to apply to training patches. If None, default transforms
         are applied.
     train_target_data : Optional[Union[str, Path]], optional
@@ -584,7 +609,7 @@ class TrainingDataWrapper(CAREamicsTrainData):
         axes: str,
         batch_size: int,
         val_data: Optional[Union[str, Path]] = None,
-        transforms: Optional[Union[List[TRANSFORMS_UNION], Compose]] = None,
+        transforms: Optional[List[TRANSFORMS_UNION]] = None,
         train_target_data: Optional[Union[str, Path]] = None,
         val_target_data: Optional[Union[str, Path]] = None,
         read_source_func: Optional[Callable] = None,
@@ -617,8 +642,8 @@ class TrainingDataWrapper(CAREamicsTrainData):
         In particular, N2V requires a specific transformation (N2V manipulates), which
         is not compatible with supervised training. The default transformations applied
         to the training patches are defined in `careamics.config.data_model`. To use
-        different transformations, pass a list of transforms or an albumentation
-        `Compose` as `transforms` parameter. See examples for more details.
+        different transformations, pass a list of transforms. See examples for more
+        details.
 
         By default, CAREamics only supports types defined in
         `careamics.config.support.SupportedData`. To read custom data types, you can set
@@ -655,7 +680,7 @@ class TrainingDataWrapper(CAREamicsTrainData):
             Batch size.
         val_data : Optional[Union[str, Path]], optional
             Validation data, by default None.
-        transforms : Optional[Union[List[TRANSFORMS_UNION], Compose]], optional
+        transforms : Optional[List[TRANSFORMS_UNION]], optional
             List of transforms to apply to training patches. If None, default transforms
             are applied.
         train_target_data : Optional[Union[str, Path]], optional
@@ -709,10 +734,7 @@ class TrainingDataWrapper(CAREamicsTrainData):
         self.data_config = DataConfig(**data_dict)
 
         # N2V specific checks, N2V, structN2V, and transforms
-        if (
-            self.data_config.has_transform_list()
-            and self.data_config.has_n2v_manipulate()
-        ):
+        if self.data_config.has_n2v_manipulate():
             # there is not target, n2v2 and structN2V can be changed
             if train_target_data is None:
                 self.data_config.set_N2V2(use_n2v2)

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.
@@ -162,7 +168,7 @@ class CAREamicsModule(L.LightningModule):
             mean=self._trainer.datamodule.predict_dataset.mean,
             std=self._trainer.datamodule.predict_dataset.std,
         )
-        denormalized_output = denorm(image=output)["image"]
+        denormalized_output, _ = denorm(patch=output)
 
         if len(aux) > 0:
             return denormalized_output, aux

careamics/lightning_prediction_datamodule.py

@@ -1,10 +1,10 @@
 """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
-from albumentations import Compose
 from torch.utils.data import DataLoader
 from torch.utils.data.dataloader import default_collate
 
@@ -39,7 +39,7 @@ def _collate_tiles(batch: List[Tuple[np.ndarray, TileInformation]]) -> Any:
 
     Parameters
     ----------
-    batch : Tuple[Tuple[np.ndarray, TileInformation], ...]
+    batch : List[Tuple[np.ndarray, TileInformation], ...]
         Batch of tiles.
 
     Returns
@@ -257,14 +257,13 @@ class PredictDataWrapper(CAREamicsPredictData):
 
     The default transformations applied to the images are defined in
     `careamics.config.inference_model`. To use different transformations, pass a list
-    of transforms or an albumentation `Compose` as `transforms` parameter. See examples
+    of transforms. See examples
     for more details.
 
     The `mean` and `std` parameters are only used if Normalization is defined either
-    in the default transformations or in the `transforms` parameter, but not with
-    a `Compose` object. If you pass a `Normalization` transform in a list as
-    `transforms`, then the mean and std parameters will be overwritten by those passed
-    to this method.
+    in the default transformations or in the `transforms` parameter. If you pass a
+    `Normalization` transform in a list as `transforms`, then the mean and std
+    parameters will be overwritten by those passed to this method.
 
     By default, CAREamics only supports types defined in
     `careamics.config.support.SupportedData`. To read custom data types, you can set
@@ -276,6 +275,12 @@ class PredictDataWrapper(CAREamicsPredictData):
     dataloaders, except for `batch_size`, which is set by the `batch_size`
     parameter.
 
+    Note that if you are using a UNet model and tiling, the tile size must be
+    divisible in every dimension by 2**d, where d is the depth of the model. This
+    avoids artefacts arising from the broken shift invariance induced by the
+    pooling layers of the UNet. If your image has less dimensions, as it may
+    happen in the Z dimension, consider padding your image.
+
     Parameters
     ----------
     pred_data : Union[str, Path, np.ndarray]
@@ -298,9 +303,6 @@ class PredictDataWrapper(CAREamicsPredictData):
         Batch size.
     tta_transforms : bool, optional
         Use test time augmentation, by default True.
-    transforms : Optional[Union[List[TRANSFORMS_UNION], Compose]], 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.
@@ -321,7 +323,6 @@ class PredictDataWrapper(CAREamicsPredictData):
         axes: str = "YX",
         batch_size: int = 1,
         tta_transforms: bool = True,
-        transforms: Optional[Union[List, Compose]] = None,
         read_source_func: Optional[Callable] = None,
         extension_filter: str = "",
         dataloader_params: Optional[dict] = None,
@@ -351,9 +352,6 @@ class PredictDataWrapper(CAREamicsPredictData):
             Batch size.
         tta_transforms : bool, optional
             Use test time augmentation, by default True.
-        transforms : Optional[Union[List[TRANSFORMS_UNION], Compose]], 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.
@@ -364,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,
@@ -373,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,7 +1,5 @@
 """Losses module."""
 
+__all__ = ["loss_factory"]
 
 from .loss_factory import loss_factory
-
-# from .noise_model_factory import noise_model_factory as noise_model_factory
-# from .noise_models import GaussianMixtureNoiseModel, HistogramNoiseModel

careamics/losses/loss_factory.py

@@ -3,6 +3,7 @@ Loss factory module.
 
 This module contains a factory function for creating loss functions.
 """
+
 from typing import Callable, Union
 
 from ..config.support import SupportedLoss
@@ -16,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/__init__.py

@@ -1,6 +1,5 @@
 """Model I/O utilities."""
 
-
 __all__ = ["load_pretrained", "export_to_bmz"]
 
 

careamics/model_io/bioimage/_readme_factory.py

@@ -1,4 +1,5 @@
 """Functions used to create a README.md file for BMZ export."""
+
 from pathlib import Path
 from typing import Optional
 
@@ -117,4 +118,4 @@ def readme_factory(
 
         readme.write_text("".join(description))
 
-    return readme
+        return readme.absolute()

careamics/model_io/bioimage/bioimage_utils.py

@@ -1,4 +1,5 @@
 """Bioimage.io utils."""
+
 from pathlib import Path
 from typing import Union
 

careamics/model_io/bioimage/model_description.py

@@ -1,4 +1,5 @@
 """Module use to build BMZ model description."""
+
 from pathlib import Path
 from typing import List, Optional, Tuple, Union
 

careamics/model_io/bmz_io.py

@@ -1,4 +1,5 @@
 """Function to export to the BioImage Model Zoo format."""
+
 import tempfile
 from pathlib import Path
 from typing import List, Optional, Tuple, Union
@@ -103,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
@@ -177,7 +178,7 @@ def export_to_bmz(
         )
 
         # test model description
-        summary: ValidationSummary = test_model(model_description)
+        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