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

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

@@ -1,10 +1,11 @@
 """Prediction Lightning data modules."""
 
 from pathlib import Path
-from typing import Any, Callable, Dict, 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 careamics.config import InferenceConfig
@@ -15,11 +16,9 @@ from careamics.dataset import (
     IterablePredDataset,
     IterableTiledPredDataset,
 )
-from careamics.dataset.dataset_utils import (
-    get_read_func,
-    list_files,
-)
+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[
@@ -32,7 +31,7 @@ PredictDatasetType = Union[
 logger = get_logger(__name__)
 
 
-class CAREamicsPredictData(L.LightningDataModule):
+class PredictDataModule(L.LightningDataModule):
     """
     CAREamics Lightning prediction data module.
 
@@ -51,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 "".
@@ -64,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,
@@ -87,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 "".
@@ -222,33 +221,36 @@ class CAREamicsPredictData(L.LightningDataModule):
             batch_size=self.batch_size,
             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
@@ -259,113 +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.
+    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[int, ...]
+    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],
-        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,
-        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.
-        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]
-            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,
-            "image_means": image_means,
-            "image_stds": image_stds,
-            "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,
+    )