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

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:
@@ -53,13 +55,4 @@ def read_tiff(file_path: Path, *args: list, **kwargs: dict) -> np.ndarray:
     else:
         raise ValueError(f"File {file_path} is not a valid tiff.")
 
-    # check dimensions
-    # TODO or should this really be done here? probably in the LightningDataModule
-    # TODO this should also be centralized somewhere else (validate_dimensions)
-    if len(array.shape) < 2 or len(array.shape) > 6:
-        raise ValueError(
-            f"Incorrect data dimensions. Must be 2, 3 or 4 (got {array.shape} for"
-            f"file {file_path})."
-        )
-
     return array

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.
@@ -148,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)
@@ -165,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
 
@@ -198,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))