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

careamics/dataset/patching/patching.py

@@ -2,9 +2,10 @@
 
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Callable, List, Tuple, Union
+from typing import Callable, Union
 
 import numpy as np
+from numpy.typing import NDArray
 
 from ...utils.logging import get_logger
 from ..dataset_utils import reshape_array
@@ -18,34 +19,49 @@ logger = get_logger(__name__)
 class Stats:
     """Dataclass to store statistics."""
 
-    means: Union[np.ndarray, tuple, list, None]
-    stds: Union[np.ndarray, tuple, list, None]
+    means: Union[NDArray, tuple, list, None]
+    """Mean of the data across channels."""
 
+    stds: Union[NDArray, tuple, list, None]
+    """Standard deviation of the data across channels."""
 
-@dataclass
-class PatchedOutput:
-    """Dataclass to store patches and statistics."""
+    def get_statistics(self) -> tuple[list[float], list[float]]:
+        """Return the means and standard deviations.
 
-    patches: Union[np.ndarray]
-    targets: Union[np.ndarray, None]
-    image_stats: Stats
-    target_stats: Stats
+        Returns
+        -------
+        tuple of two lists of floats
+            Means and standard deviations.
+        """
+        if self.means is None or self.stds is None:
+            return [], []
+
+        return list(self.means), list(self.stds)
 
 
 @dataclass
-class StatsOutput:
+class PatchedOutput:
     """Dataclass to store patches and statistics."""
 
+    patches: Union[NDArray]
+    """Image patches."""
+
+    targets: Union[NDArray, None]
+    """Target patches."""
+
     image_stats: Stats
+    """Statistics of the image patches."""
+
     target_stats: Stats
+    """Statistics of the target patches."""
 
 
 # called by in memory dataset
 def prepare_patches_supervised(
-    train_files: List[Path],
-    target_files: List[Path],
+    train_files: list[Path],
+    target_files: list[Path],
     axes: str,
-    patch_size: Union[List[int], Tuple[int, ...]],
+    patch_size: Union[list[int], tuple[int, ...]],
     read_source_func: Callable,
 ) -> PatchedOutput:
     """
@@ -55,13 +71,13 @@ def prepare_patches_supervised(
 
     Parameters
     ----------
-    train_files : List[Path]
+    train_files : list of pathlib.Path
         List of paths to training data.
-    target_files : List[Path]
+    target_files : list of pathlib.Path
         List of paths to target data.
     axes : str
         Axes of the data.
-    patch_size : Union[List[int], Tuple[int]]
+    patch_size : list or tuple of int
         Size of the patches.
     read_source_func : Callable
         Function to read the data.
@@ -127,9 +143,9 @@ def prepare_patches_supervised(
 
 # called by in_memory_dataset
 def prepare_patches_unsupervised(
-    train_files: List[Path],
+    train_files: list[Path],
     axes: str,
-    patch_size: Union[List[int], Tuple[int]],
+    patch_size: Union[list[int], tuple[int]],
     read_source_func: Callable,
 ) -> PatchedOutput:
     """Iterate over data source and create an array of patches.
@@ -138,19 +154,19 @@ def prepare_patches_unsupervised(
 
     Parameters
     ----------
-    train_files : List[Path]
+    train_files : list of pathlib.Path
         List of paths to training data.
     axes : str
         Axes of the data.
-    patch_size : Union[List[int], Tuple[int]]
+    patch_size : list or tuple of int
         Size of the patches.
     read_source_func : Callable
         Function to read the data.
 
     Returns
     -------
-    Tuple[np.ndarray, None, float, float]
-        Source and target patches, mean and standard deviation.
+    PatchedOutput
+        Dataclass holding patches and their statistics.
     """
     means, stds, num_samples = 0, 0, 0
     all_patches = []
@@ -189,10 +205,10 @@ def prepare_patches_unsupervised(
 
 # called on arrays by in memory dataset
 def prepare_patches_supervised_array(
-    data: np.ndarray,
+    data: NDArray,
     axes: str,
-    data_target: np.ndarray,
-    patch_size: Union[List[int], Tuple[int]],
+    data_target: NDArray,
+    patch_size: Union[list[int], tuple[int]],
 ) -> PatchedOutput:
     """Iterate over data source and create an array of patches.
 
@@ -203,19 +219,19 @@ def prepare_patches_supervised_array(
 
     Parameters
     ----------
-    data : np.ndarray
+    data : numpy.ndarray
         Input data array.
     axes : str
         Axes of the data.
-    data_target : np.ndarray
+    data_target : numpy.ndarray
         Target data array.
-    patch_size : Union[List[int], Tuple[int]]
+    patch_size : list or tuple of int
         Size of the patches.
 
     Returns
     -------
-    Tuple[np.ndarray, np.ndarray, float, float]
-        Source and target patches, mean and standard deviation.
+    PatchedOutput
+        Dataclass holding the source and target patches, with their statistics.
     """
     # reshape array
     reshaped_sample = reshape_array(data, axes)
@@ -245,9 +261,9 @@ def prepare_patches_supervised_array(
 
 # called by in memory dataset
 def prepare_patches_unsupervised_array(
-    data: np.ndarray,
+    data: NDArray,
     axes: str,
-    patch_size: Union[List[int], Tuple[int]],
+    patch_size: Union[list[int], tuple[int]],
 ) -> PatchedOutput:
     """
     Iterate over data source and create an array of patches.
@@ -259,17 +275,17 @@ def prepare_patches_unsupervised_array(
 
     Parameters
     ----------
-    data : np.ndarray
+    data : numpy.ndarray
         Input data array.
     axes : str
         Axes of the data.
-    patch_size : Union[List[int], Tuple[int]]
+    patch_size : list or tuple of int
         Size of the patches.
 
     Returns
     -------
-    Tuple[np.ndarray, None, float, float]
-        Source patches, mean and standard deviation.
+    PatchedOutput
+        Dataclass holding the patches and their statistics.
     """
     # reshape array
     reshaped_sample = reshape_array(data, axes)

careamics/file_io/__init__.py

@@ -0,0 +1,7 @@
+"""Functions relating reading and writing image files."""
+
+__all__ = ["read", "write", "get_read_func", "get_write_func"]
+
+from . import read, write
+from .read import get_read_func
+from .write import get_write_func

careamics/file_io/read/__init__.py

@@ -0,0 +1,11 @@
+"""Functions relating to reading image files of different formats."""
+
+__all__ = [
+    "get_read_func",
+    "read_tiff",
+    "read_zarr",
+]
+
+from .get_func import get_read_func
+from .tiff import read_tiff
+from .zarr import read_zarr

careamics/file_io/read/get_func.py

@@ -0,0 +1,56 @@
+"""Module to get read functions."""
+
+from pathlib import Path
+from typing import Callable, Dict, Protocol, Union
+
+from numpy.typing import NDArray
+
+from careamics.config.support import SupportedData
+
+from .tiff import read_tiff
+
+
+# This is very strict, function signature has to match including arg names
+# See WriteFunc notes
+class ReadFunc(Protocol):
+    """Protocol for type hinting read functions."""
+
+    def __call__(self, file_path: Path, *args, **kwargs) -> NDArray:
+        """
+        Type hinted callables must match this function signature (not including self).
+
+        Parameters
+        ----------
+        file_path : pathlib.Path
+            Path to file.
+        *args
+            Other positional arguments.
+        **kwargs
+            Other keyword arguments.
+        """
+
+
+READ_FUNCS: Dict[SupportedData, ReadFunc] = {
+    SupportedData.TIFF: read_tiff,
+}
+
+
+def get_read_func(data_type: Union[str, SupportedData]) -> Callable:
+    """
+    Get the read function for the data type.
+
+    Parameters
+    ----------
+    data_type : SupportedData
+        Data type.
+
+    Returns
+    -------
+    callable
+        Read function.
+    """
+    if data_type in READ_FUNCS:
+        data_type = SupportedData(data_type)  # mypy complaining about dict key type
+        return READ_FUNCS[data_type]
+    else:
+        raise NotImplementedError(f"Data type '{data_type}' is not supported.")

careamics/{dataset/dataset_utils/read_tiff.py → file_io/read/tiff.py}

@@ -44,7 +44,9 @@ def read_tiff(file_path: Path, *args: list, **kwargs: dict) -> np.ndarray:
     ValueError
         If the axes length is incorrect.
     """
-    if fnmatch(file_path.suffix, SupportedData.get_extension(SupportedData.TIFF)):
+    if fnmatch(
+        file_path.suffix, SupportedData.get_extension_pattern(SupportedData.TIFF)
+    ):
         try:
             array = tifffile.imread(file_path)
         except (ValueError, OSError) as e:

careamics/file_io/write/__init__.py

@@ -0,0 +1,9 @@
+"""Functions relating to writing image files of different formats."""
+
+__all__ = [
+    "get_write_func",
+    "write_tiff",
+]
+
+from .get_func import get_write_func
+from .tiff import write_tiff

careamics/file_io/write/get_func.py

@@ -0,0 +1,59 @@
+"""Module to get write functions."""
+
+from pathlib import Path
+from typing import Protocol, Union
+
+from numpy.typing import NDArray
+
+from careamics.config.support import SupportedData
+
+from .tiff import write_tiff
+
+
+# This is very strict, arguments have to be called file_path & img
+# Alternative? - doesn't capture *args & **kwargs
+# WriteFunc = Callable[[Path, NDArray], None]
+class WriteFunc(Protocol):
+    """Protocol for type hinting write functions."""
+
+    def __call__(self, file_path: Path, img: NDArray, *args, **kwargs) -> None:
+        """
+        Type hinted callables must match this function signature (not including self).
+
+        Parameters
+        ----------
+        file_path : pathlib.Path
+            Path to file.
+        img : numpy.ndarray
+            Image data to save.
+        *args
+            Other positional arguments.
+        **kwargs
+            Other keyword arguments.
+        """
+
+
+WRITE_FUNCS: dict[SupportedData, WriteFunc] = {
+    SupportedData.TIFF: write_tiff,
+}
+
+
+def get_write_func(data_type: Union[str, SupportedData]) -> WriteFunc:
+    """
+    Get the write function for the data type.
+
+    Parameters
+    ----------
+    data_type : SupportedData
+        Data type.
+
+    Returns
+    -------
+    callable
+        Write function.
+    """
+    if data_type in WRITE_FUNCS:
+        data_type = SupportedData(data_type)  # mypy complaining about dict key type
+        return WRITE_FUNCS[data_type]
+    else:
+        raise NotImplementedError(f"Data type {data_type} is not supported.")

careamics/file_io/write/tiff.py

@@ -0,0 +1,39 @@
+"""Write tiff function."""
+
+from fnmatch import fnmatch
+from pathlib import Path
+
+import tifffile
+from numpy.typing import NDArray
+
+from careamics.config.support import SupportedData
+
+
+def write_tiff(file_path: Path, img: NDArray, *args, **kwargs) -> None:
+    """
+    Write tiff files.
+
+    Parameters
+    ----------
+    file_path : pathlib.Path
+        Path to file.
+    img : numpy.ndarray
+        Image data to save.
+    *args
+        Positional arguments passed to `tifffile.imwrite`.
+    **kwargs
+        Keyword arguments passed to `tifffile.imwrite`.
+
+    Raises
+    ------
+    ValueError
+        When the file extension of `file_path` does not match the Unix shell-style
+        pattern '*.tif*'.
+    """
+    if not fnmatch(
+        file_path.suffix, SupportedData.get_extension_pattern(SupportedData.TIFF)
+    ):
+        raise ValueError(
+            f"Unexpected extension '{file_path.suffix}' for save file type 'tiff'."
+        )
+    tifffile.imwrite(file_path, img, *args, **kwargs)

careamics/lightning/__init__.py

@@ -0,0 +1,17 @@
+"""CAREamics PyTorch Lightning modules."""
+
+__all__ = [
+    "CAREamicsModule",
+    "create_careamics_module",
+    "TrainDataModule",
+    "create_train_datamodule",
+    "PredictDataModule",
+    "create_predict_datamodule",
+    "HyperParametersCallback",
+    "ProgressBarCallback",
+]
+
+from .callbacks import HyperParametersCallback, ProgressBarCallback
+from .lightning_module import CAREamicsModule, create_careamics_module
+from .predict_data_module import PredictDataModule, create_predict_datamodule
+from .train_data_module import TrainDataModule, create_train_datamodule

careamics/{lightning_module.py → lightning/lightning_module.py}

@@ -23,19 +23,19 @@ class CAREamicsModule(L.LightningModule):
     """
     CAREamics Lightning module.
 
-    This class encapsulates the a PyTorch model along with the training, validation,
+    This class encapsulates the 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_config : AlgorithmModel or dict
         Algorithm configuration.
 
     Attributes
     ----------
-    model : nn.Module
+    model : torch.nn.Module
         PyTorch model.
-    loss_func : nn.Module
+    loss_func : torch.nn.Module
         Loss function.
     optimizer_name : str
         Optimizer name.
@@ -53,7 +53,7 @@ class CAREamicsModule(L.LightningModule):
 
         Parameters
         ----------
-        algorithm_config : Union[AlgorithmModel, dict]
+        algorithm_config : AlgorithmModel or dict
             Algorithm configuration.
         """
         super().__init__()
@@ -91,7 +91,7 @@ class CAREamicsModule(L.LightningModule):
 
         Parameters
         ----------
-        batch : Tensor
+        batch : torch.Tensor
             Input batch.
         batch_idx : Any
             Batch index.
@@ -114,7 +114,7 @@ class CAREamicsModule(L.LightningModule):
 
         Parameters
         ----------
-        batch : Tensor
+        batch : torch.Tensor
             Input batch.
         batch_idx : Any
             Batch index.
@@ -138,7 +138,7 @@ class CAREamicsModule(L.LightningModule):
 
         Parameters
         ----------
-        batch : Tensor
+        batch : torch.Tensor
             Input batch.
         batch_idx : Any
             Batch index.
@@ -202,101 +202,74 @@ class CAREamicsModule(L.LightningModule):
         }
 
 
-class CAREamicsModuleWrapper(CAREamicsModule):
-    """Class defining the API for CAREamics Lightning layer.
+def create_careamics_module(
+    algorithm: Union[SupportedAlgorithm, str],
+    loss: Union[SupportedLoss, str],
+    architecture: Union[SupportedArchitecture, str],
+    model_parameters: Optional[dict] = None,
+    optimizer: Union[SupportedOptimizer, str] = "Adam",
+    optimizer_parameters: Optional[dict] = None,
+    lr_scheduler: Union[SupportedScheduler, str] = "ReduceLROnPlateau",
+    lr_scheduler_parameters: Optional[dict] = None,
+) -> CAREamicsModule:
+    """Create a CAREamics Lithgning module.
 
-    This class exposes parameters used to create an AlgorithmModel instance, triggering
-    parameters validation.
+    This function exposes parameters used to create an AlgorithmModel instance,
+    triggering parameters validation.
 
     Parameters
     ----------
-    algorithm : Union[SupportedAlgorithm, str]
+    algorithm : SupportedAlgorithm or str
         Algorithm to use for training (see SupportedAlgorithm).
-    loss : Union[SupportedLoss, str]
+    loss : SupportedLoss or str
         Loss function to use for training (see SupportedLoss).
-    architecture : Union[SupportedArchitecture, str]
+    architecture : SupportedArchitecture or str
         Model architecture to use for training (see SupportedArchitecture).
     model_parameters : dict, optional
         Model parameters to use for training, by default {}. Model parameters are
         defined in the relevant `torch.nn.Module` class, or Pyddantic model (see
         `careamics.config.architectures`).
-    optimizer : Union[SupportedOptimizer, str], optional
+    optimizer : SupportedOptimizer or str, optional
         Optimizer to use for training, by default "Adam" (see SupportedOptimizer).
     optimizer_parameters : dict, optional
         Optimizer parameters to use for training, as defined in `torch.optim`, by
         default {}.
-    lr_scheduler : Union[SupportedScheduler, str], optional
+    lr_scheduler : SupportedScheduler or str, optional
         Learning rate scheduler to use for training, by default "ReduceLROnPlateau"
         (see SupportedScheduler).
     lr_scheduler_parameters : dict, optional
         Learning rate scheduler parameters to use for training, as defined in
         `torch.optim`, by default {}.
-    """
-
-    def __init__(
-        self,
-        algorithm: Union[SupportedAlgorithm, str],
-        loss: Union[SupportedLoss, str],
-        architecture: Union[SupportedArchitecture, str],
-        model_parameters: Optional[dict] = None,
-        optimizer: Union[SupportedOptimizer, str] = "Adam",
-        optimizer_parameters: Optional[dict] = None,
-        lr_scheduler: Union[SupportedScheduler, str] = "ReduceLROnPlateau",
-        lr_scheduler_parameters: Optional[dict] = None,
-    ) -> None:
-        """
-        Wrapper for the CAREamics model, exposing all algorithm configuration arguments.
 
-        Parameters
-        ----------
-        algorithm : Union[SupportedAlgorithm, str]
-            Algorithm to use for training (see SupportedAlgorithm).
-        loss : Union[SupportedLoss, str]
-            Loss function to use for training (see SupportedLoss).
-        architecture : Union[SupportedArchitecture, str]
-            Model architecture to use for training (see SupportedArchitecture).
-        model_parameters : dict, optional
-            Model parameters to use for training, by default {}. Model parameters are
-            defined in the relevant `torch.nn.Module` class, or Pyddantic model (see
-            `careamics.config.architectures`).
-        optimizer : Union[SupportedOptimizer, str], optional
-            Optimizer to use for training, by default "Adam" (see SupportedOptimizer).
-        optimizer_parameters : dict, optional
-            Optimizer parameters to use for training, as defined in `torch.optim`, by
-            default {}.
-        lr_scheduler : Union[SupportedScheduler, str], optional
-            Learning rate scheduler to use for training, by default "ReduceLROnPlateau"
-            (see SupportedScheduler).
-        lr_scheduler_parameters : dict, optional
-            Learning rate scheduler parameters to use for training, as defined in
-            `torch.optim`, by default {}.
-        """
-        # create a AlgorithmModel compatible dictionary
-        if lr_scheduler_parameters is None:
-            lr_scheduler_parameters = {}
-        if optimizer_parameters is None:
-            optimizer_parameters = {}
-        if model_parameters is None:
-            model_parameters = {}
-        algorithm_configuration = {
-            "algorithm": algorithm,
-            "loss": loss,
-            "optimizer": {
-                "name": optimizer,
-                "parameters": optimizer_parameters,
-            },
-            "lr_scheduler": {
-                "name": lr_scheduler,
-                "parameters": lr_scheduler_parameters,
-            },
-        }
-        model_configuration = {"architecture": architecture}
-        model_configuration.update(model_parameters)
-
-        # add model parameters to algorithm configuration
-        algorithm_configuration["model"] = model_configuration
-
-        # call the parent init using an AlgorithmModel instance
-        super().__init__(AlgorithmConfig(**algorithm_configuration))
-
-        # TODO add load_from_checkpoint wrapper
+    Returns
+    -------
+    CAREamicsModule
+        CAREamics Lightning module.
+    """
+    # create a AlgorithmModel compatible dictionary
+    if lr_scheduler_parameters is None:
+        lr_scheduler_parameters = {}
+    if optimizer_parameters is None:
+        optimizer_parameters = {}
+    if model_parameters is None:
+        model_parameters = {}
+    algorithm_configuration = {
+        "algorithm": algorithm,
+        "loss": loss,
+        "optimizer": {
+            "name": optimizer,
+            "parameters": optimizer_parameters,
+        },
+        "lr_scheduler": {
+            "name": lr_scheduler,
+            "parameters": lr_scheduler_parameters,
+        },
+    }
+    model_configuration = {"architecture": architecture}
+    model_configuration.update(model_parameters)
+
+    # add model parameters to algorithm configuration
+    algorithm_configuration["model"] = model_configuration
+
+    # call the parent init using an AlgorithmModel instance
+    return CAREamicsModule(AlgorithmConfig(**algorithm_configuration))