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

careamics/{lightning_prediction_datamodule.py → lightning/predict_data_module.py}

@@ -1,69 +1,37 @@
 """Prediction Lightning data modules."""
 
 from pathlib import Path
-from typing import Any, Callable, Dict, List, Literal, Optional, Tuple, Union
+from typing import Any, Callable, Literal, Optional, Union
 
 import numpy as np
 import pytorch_lightning as L
+from numpy.typing import NDArray
 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.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 import (
+    InMemoryPredDataset,
+    InMemoryTiledPredDataset,
+    IterablePredDataset,
+    IterableTiledPredDataset,
 )
+from careamics.dataset.dataset_utils import list_files
+from careamics.dataset.tiling.collate_tiles import collate_tiles
+from careamics.file_io.read import get_read_func
 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):
+class PredictDataModule(L.LightningDataModule):
     """
     CAREamics Lightning prediction data module.
 
@@ -82,9 +50,9 @@ class CAREamicsPredictData(L.LightningDataModule):
     ----------
     pred_config : InferenceModel
         Pydantic model for CAREamics prediction configuration.
-    pred_data : Union[Path, str, np.ndarray]
+    pred_data : pathlib.Path or str or numpy.ndarray
         Prediction data, can be a path to a folder, a file or a numpy array.
-    read_source_func : Optional[Callable], optional
+    read_source_func : Callable, optional
         Function to read custom types, by default None.
     extension_filter : str, optional
         Filter to filter file extensions for custom types, by default "".
@@ -95,7 +63,7 @@ class CAREamicsPredictData(L.LightningDataModule):
     def __init__(
         self,
         pred_config: InferenceConfig,
-        pred_data: Union[Path, str, np.ndarray],
+        pred_data: Union[Path, str, NDArray],
         read_source_func: Optional[Callable] = None,
         extension_filter: str = "",
         dataloader_params: Optional[dict] = None,
@@ -118,9 +86,9 @@ class CAREamicsPredictData(L.LightningDataModule):
         ----------
         pred_config : InferenceModel
             Pydantic model for CAREamics prediction configuration.
-        pred_data : Union[Path, str, np.ndarray]
+        pred_data : pathlib.Path or str or numpy.ndarray
             Prediction data, can be a path to a folder, a file or a numpy array.
-        read_source_func : Optional[Callable], optional
+        read_source_func : Callable, optional
             Function to read custom types, by default None.
         extension_filter : str, optional
             Filter to filter file extensions for custom types, by default "".
@@ -182,6 +150,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 +183,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,35 +219,38 @@ 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
-
+        )
 
-class PredictDataWrapper(CAREamicsPredictData):
-    """
-    Wrapper around the CAREamics inference Lightning data module.
 
-    This class is used to explicitely pass the parameters usually contained in a
+def create_predict_datamodule(
+    pred_data: Union[str, Path, NDArray],
+    data_type: Union[Literal["array", "tiff", "custom"], SupportedData],
+    axes: str,
+    image_means: list[float],
+    image_stds: list[float],
+    tile_size: Optional[tuple[int, ...]] = None,
+    tile_overlap: Optional[tuple[int, ...]] = None,
+    batch_size: int = 1,
+    tta_transforms: bool = True,
+    read_source_func: Optional[Callable] = None,
+    extension_filter: str = "",
+    dataloader_params: Optional[dict] = None,
+) -> PredictDataModule:
+    """Create a CAREamics prediction Lightning datamodule.
+
+    This function is used to explicitely pass the parameters usually contained in an
     `inference_model` configuration.
 
     Since the lightning datamodule has no access to the model, make sure that the
     parameters passed to the datamodule are consistent with the model's requirements
-    and are coherent.
+    and are coherent. This can be done by creating a `Configuration` object beforehand
+    and passing its parameters to the different Lightning modules.
 
     The data module can be used with Path, str or numpy arrays. To use array data, set
     `data_type` to `array` and pass a numpy array to `train_data`.
 
-    The default transformations applied to the images are defined in
-    `careamics.config.inference_model`. To use different transformations, pass a list
-    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. 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
     `data_type` to `custom` and provide a function that returns a numpy array from a
@@ -275,117 +261,73 @@ 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]
+    pred_data : str or pathlib.Path or numpy.ndarray
         Prediction data.
-    data_type : Union[Literal["array", "tiff", "custom"], SupportedData]
+    data_type : {"array", "tiff", "custom"}
         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.
-    tile_size : Tuple[int, ...]
+    axes : str
+        Axes of the data, choosen among SCZYX.
+    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 of int, optional
         Tile size, 2D or 3D tile size.
-    tile_overlap : Tuple[int, ...]
+    tile_overlap : tuple of int, optional
         Tile overlap, 2D or 3D tile overlap.
-    axes : str
-        Axes of the data, choosen amongst SCZYX.
     batch_size : int
         Batch size.
     tta_transforms : bool, optional
         Use test time augmentation, by default True.
-    read_source_func : Optional[Callable], optional
+    read_source_func : Callable, optional
         Function to read the source data, used if `data_type` is `custom`, by
         default None.
     extension_filter : str, optional
         Filter for file extensions, used if `data_type` is `custom`, by default "".
     dataloader_params : dict, optional
         Pytorch dataloader parameters, by default {}.
-    """
 
-    def __init__(
-        self,
-        pred_data: Union[str, Path, np.ndarray],
-        data_type: Union[Literal["array", "tiff", "custom"], SupportedData],
-        mean: float,
-        std: float,
-        tile_size: Optional[Tuple[int, ...]] = None,
-        tile_overlap: Optional[Tuple[int, ...]] = None,
-        axes: str = "YX",
-        batch_size: int = 1,
-        tta_transforms: bool = True,
-        read_source_func: Optional[Callable] = None,
-        extension_filter: str = "",
-        dataloader_params: Optional[dict] = None,
-    ) -> None:
-        """
-        Constructor.
+    Returns
+    -------
+    PredictDataModule
+        CAREamics prediction datamodule.
 
-        Parameters
-        ----------
-        pred_data : Union[str, Path, np.ndarray]
-            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.
-        tile_size : List[int]
-            Tile size, 2D or 3D tile size.
-        tile_overlap : List[int]
-            Tile overlap, 2D or 3D tile overlap.
-        axes : str
-            Axes of the data, choosen amongst SCZYX.
-        batch_size : int
-            Batch size.
-        tta_transforms : bool, optional
-            Use test time augmentation, by default True.
-        read_source_func : Optional[Callable], optional
-            Function to read the source data, used if `data_type` is `custom`, by
-            default None.
-        extension_filter : str, optional
-            Filter for file extensions, used if `data_type` is `custom`, by default "".
-        dataloader_params : dict, optional
-            Pytorch dataloader parameters, by default {}.
-        """
-        if dataloader_params is None:
-            dataloader_params = {}
-        prediction_dict: Dict[str, Any] = {
-            "data_type": data_type,
-            "tile_size": tile_size,
-            "tile_overlap": tile_overlap,
-            "axes": axes,
-            "mean": mean,
-            "std": std,
-            "tta": tta_transforms,
-            "batch_size": batch_size,
-            "transforms": [],
-        }
-
-        # validate configuration
-        self.prediction_config = InferenceConfig(**prediction_dict)
-
-        # sanity check on the dataloader parameters
-        if "batch_size" in dataloader_params:
-            # remove it
-            del dataloader_params["batch_size"]
-
-        super().__init__(
-            pred_config=self.prediction_config,
-            pred_data=pred_data,
-            read_source_func=read_source_func,
-            extension_filter=extension_filter,
-            dataloader_params=dataloader_params,
-        )
+    Notes
+    -----
+    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.
+    """
+    if dataloader_params is None:
+        dataloader_params = {}
+
+    prediction_dict: dict[str, Any] = {
+        "data_type": data_type,
+        "tile_size": tile_size,
+        "tile_overlap": tile_overlap,
+        "axes": axes,
+        "image_means": image_means,
+        "image_stds": image_stds,
+        "tta_transforms": tta_transforms,
+        "batch_size": batch_size,
+    }
+
+    # validate configuration
+    prediction_config = InferenceConfig(**prediction_dict)
+
+    # sanity check on the dataloader parameters
+    if "batch_size" in dataloader_params:
+        # remove it
+        del dataloader_params["batch_size"]
+
+    return PredictDataModule(
+        pred_config=prediction_config,
+        pred_data=pred_data,
+        read_source_func=read_source_func,
+        extension_filter=extension_filter,
+        dataloader_params=dataloader_params,
+    )