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

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
+        )
 
-        # data
-        self.train_data = train_data
-        self.val_data = val_data
+        self.train_data_target: Union[Path, np.ndarray] = (
+            Path(train_data_target)
+            if isinstance(train_data_target, str)
+            else train_data_target
+        )
 
-        self.train_data_target = train_data_target
-        self.val_data_target = val_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,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,
@@ -557,12 +583,7 @@ class TrainingDataWrapper(CAREamicsTrainData):
     >>> my_array = np.arange(256).reshape(16, 16)
     >>> my_transforms = [
     ...     {
-    ...         "name": SupportedTransform.NORMALIZE.value,
-    ...         "mean": 0,
-    ...         "std": 1,
-    ...     },
-    ...     {
-    ...         "name": SupportedTransform.N2V_MANIPULATE.value,
+    ...         "name": SupportedTransform.XY_FLIP.value,
     ...     }
     ... ]
     >>> data_module = TrainingDataWrapper(

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.
@@ -142,13 +148,17 @@ class CAREamicsModule(L.LightningModule):
         Any
             Model output.
         """
-        x, *aux = batch
+        if self._trainer.datamodule.tiled:
+            x, *aux = batch
+        else:
+            x = batch
+            aux = []
 
         # apply test-time augmentation if available
         # TODO: probably wont work with batch size > 1
         if self._trainer.datamodule.prediction_config.tta_transforms:
             tta = ImageRestorationTTA()
-            augmented_batch = tta.forward(batch[0])  # list of augmented tensors
+            augmented_batch = tta.forward(x)  # list of augmented tensors
             augmented_output = []
             for augmented in augmented_batch:
                 augmented_pred = self.model(augmented)
@@ -159,13 +169,13 @@ class CAREamicsModule(L.LightningModule):
 
         # Denormalize the output
         denorm = Denormalize(
-            mean=self._trainer.datamodule.predict_dataset.mean,
-            std=self._trainer.datamodule.predict_dataset.std,
+            image_means=self._trainer.datamodule.predict_dataset.image_means,
+            image_stds=self._trainer.datamodule.predict_dataset.image_stds,
         )
-        denormalized_output, _ = denorm(patch=output)
+        denormalized_output = denorm(patch=output.cpu().numpy())
 
-        if len(aux) > 0:
-            return denormalized_output, aux
+        if len(aux) > 0:  # aux can be tiling information
+            return denormalized_output, *aux
         else:
             return denormalized_output
 

careamics/lightning_prediction_datamodule.py

@@ -1,68 +1,37 @@
 """Prediction Lightning data modules."""
 
 from pathlib import Path
-from typing import Any, Callable, List, Literal, Optional, Tuple, Union
+from typing import Any, Callable, Dict, Literal, Optional, Tuple, Union
 
 import numpy as np
 import pytorch_lightning as L
 from torch.utils.data import DataLoader
-from torch.utils.data.dataloader import default_collate
 
 from careamics.config import InferenceConfig
 from careamics.config.support import SupportedData
-from careamics.config.tile_information import TileInformation
+from careamics.dataset import (
+    InMemoryPredDataset,
+    InMemoryTiledPredDataset,
+    IterablePredDataset,
+    IterableTiledPredDataset,
+)
 from careamics.dataset.dataset_utils import (
     get_read_func,
     list_files,
 )
-from careamics.dataset.in_memory_dataset import (
-    InMemoryPredictionDataset,
-)
-from careamics.dataset.iterable_dataset import (
-    IterablePredictionDataset,
-)
+from careamics.dataset.tiling.collate_tiles import collate_tiles
 from careamics.utils import get_logger
 
-PredictDatasetType = Union[InMemoryPredictionDataset, IterablePredictionDataset]
+PredictDatasetType = Union[
+    InMemoryPredDataset,
+    InMemoryTiledPredDataset,
+    IterablePredDataset,
+    IterableTiledPredDataset,
+]
 
 logger = get_logger(__name__)
 
 
-def _collate_tiles(batch: List[Tuple[np.ndarray, TileInformation]]) -> Any:
-    """
-    Collate tiles received from CAREamics prediction dataloader.
-
-    CAREamics prediction dataloader returns tuples of arrays and TileInformation. In
-    case of non-tiled data, this function will return the arrays. In case of tiled data,
-    it will return the arrays, the last tile flag, the overlap crop coordinates and the
-    stitch coordinates.
-
-    Parameters
-    ----------
-    batch : List[Tuple[np.ndarray, TileInformation], ...]
-        Batch of tiles.
-
-    Returns
-    -------
-    Any
-        Collated batch.
-    """
-    first_tile_info: TileInformation = batch[0][1]
-    # if not tiled, then return arrays
-    if not first_tile_info.tiled:
-        arrays, _ = zip(*batch)
-
-        return default_collate(arrays)
-    # else we explicit the last_tile flag and coordinates
-    else:
-        new_batch = [
-            (tile, t.last_tile, t.array_shape, t.overlap_crop_coords, t.stitch_coords)
-            for tile, t in batch
-        ]
-
-        return default_collate(new_batch)
-
-
 class CAREamicsPredictData(L.LightningDataModule):
     """
     CAREamics Lightning prediction data module.
@@ -182,6 +151,9 @@ class CAREamicsPredictData(L.LightningDataModule):
         self.tile_size = pred_config.tile_size
         self.tile_overlap = pred_config.tile_overlap
 
+        # check if it is tiled
+        self.tiled = self.tile_size is not None and self.tile_overlap is not None
+
         # read source function
         if pred_config.data_type == SupportedData.CUSTOM:
             # mypy check
@@ -212,17 +184,29 @@ class CAREamicsPredictData(L.LightningDataModule):
         """
         # if numpy array
         if self.data_type == SupportedData.ARRAY:
-            # prediction dataset
-            self.predict_dataset: PredictDatasetType = InMemoryPredictionDataset(
-                prediction_config=self.prediction_config,
-                inputs=self.pred_data,
-            )
+            if self.tiled:
+                self.predict_dataset: PredictDatasetType = InMemoryTiledPredDataset(
+                    prediction_config=self.prediction_config,
+                    inputs=self.pred_data,
+                )
+            else:
+                self.predict_dataset = InMemoryPredDataset(
+                    prediction_config=self.prediction_config,
+                    inputs=self.pred_data,
+                )
         else:
-            self.predict_dataset = IterablePredictionDataset(
-                prediction_config=self.prediction_config,
-                src_files=self.pred_files,
-                read_source_func=self.read_source_func,
-            )
+            if self.tiled:
+                self.predict_dataset = IterableTiledPredDataset(
+                    prediction_config=self.prediction_config,
+                    src_files=self.pred_files,
+                    read_source_func=self.read_source_func,
+                )
+            else:
+                self.predict_dataset = IterablePredDataset(
+                    prediction_config=self.prediction_config,
+                    src_files=self.pred_files,
+                    read_source_func=self.read_source_func,
+                )
 
     def predict_dataloader(self) -> DataLoader:
         """
@@ -236,7 +220,7 @@ class CAREamicsPredictData(L.LightningDataModule):
         return DataLoader(
             self.predict_dataset,
             batch_size=self.batch_size,
-            collate_fn=_collate_tiles,
+            collate_fn=collate_tiles if self.tiled else None,
             **self.dataloader_params,
         )  # TODO check workers are used
 
@@ -287,12 +271,10 @@ class PredictDataWrapper(CAREamicsPredictData):
         Prediction data.
     data_type : Union[Literal["array", "tiff", "custom"], SupportedData]
         Data type, see `SupportedData` for available options.
-    mean : float
-        Mean value for normalization, only used if Normalization is defined in the
-        transforms.
-    std : float
-        Standard deviation value for normalization, only used if Normalization is
-        defined in the transform.
+    image_means : list of float
+        Mean values for normalization, only used if Normalization is defined.
+    image_stds : list of float
+        Std values for normalization, only used if Normalization is defined.
     tile_size : Tuple[int, ...]
         Tile size, 2D or 3D tile size.
     tile_overlap : Tuple[int, ...]
@@ -303,9 +285,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.
@@ -319,14 +298,13 @@ class PredictDataWrapper(CAREamicsPredictData):
         self,
         pred_data: Union[str, Path, np.ndarray],
         data_type: Union[Literal["array", "tiff", "custom"], SupportedData],
-        mean: float,
-        std: float,
+        image_means=list[float],
+        image_stds=list[float],
         tile_size: Optional[Tuple[int, ...]] = None,
         tile_overlap: Optional[Tuple[int, ...]] = None,
         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,
@@ -340,12 +318,10 @@ class PredictDataWrapper(CAREamicsPredictData):
             Prediction data.
         data_type : Union[Literal["array", "tiff", "custom"], SupportedData]
             Data type, see `SupportedData` for available options.
-        mean : float
-            Mean value for normalization, only used if Normalization is defined in the
-            transforms.
-        std : float
-            Standard deviation value for normalization, only used if Normalization is
-            defined in the transform.
+        image_means : list of float
+            Mean values for normalization, only used if Normalization is defined.
+        image_stds : list of float
+            Std values for normalization, only used if Normalization is defined.
         tile_size : List[int]
             Tile size, 2D or 3D tile size.
         tile_overlap : List[int]
@@ -356,9 +332,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,21 +342,18 @@ 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,
             "axes": axes,
-            "mean": mean,
-            "std": std,
+            "image_means": image_means,
+            "image_stds": image_stds,
             "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/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.