careamics 0.0.5__py3-none-any.whl → 0.0.6__py3-none-any.whl

careamics/dataset/iterable_dataset.py

@@ -10,7 +10,7 @@ from typing import Callable, Optional
 import numpy as np
 from torch.utils.data import IterableDataset
 
-from careamics.config import DataConfig
+from careamics.config import GeneralDataConfig
 from careamics.config.transformations import NormalizeModel
 from careamics.file_io.read import read_tiff
 from careamics.transforms import Compose
@@ -49,7 +49,7 @@ class PathIterableDataset(IterableDataset):
 
     def __init__(
         self,
-        data_config: DataConfig,
+        data_config: GeneralDataConfig,
         src_files: list[Path],
         target_files: Optional[list[Path]] = None,
         read_source_func: Callable = read_tiff,
@@ -58,7 +58,7 @@ class PathIterableDataset(IterableDataset):
 
         Parameters
         ----------
-        data_config : DataConfig
+        data_config : GeneralDataConfig
             Data configuration.
         src_files : list[Path]
             List of data files.
@@ -115,7 +115,7 @@ class PathIterableDataset(IterableDataset):
                     target_stds=self.target_stats.stds,
                 )
             ]
-            + data_config.transforms
+            + list(data_config.transforms)
         )
 
     def _calculate_mean_and_std(self) -> tuple[Stats, Stats]:

careamics/dataset/iterable_pred_dataset.py

@@ -2,8 +2,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Generator
 from pathlib import Path
-from typing import Any, Callable, Generator
+from typing import Any, Callable
 
 from numpy.typing import NDArray
 from torch.utils.data import IterableDataset

careamics/dataset/iterable_tiled_pred_dataset.py

@@ -2,8 +2,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Generator
 from pathlib import Path
-from typing import Any, Callable, Generator
+from typing import Any, Callable
 
 from numpy.typing import NDArray
 from torch.utils.data import IterableDataset

careamics/dataset/patching/random_patching.py

@@ -1,6 +1,7 @@
 """Random patching utilities."""
 
-from typing import Generator, List, Optional, Tuple, Union
+from collections.abc import Generator
+from typing import Optional, Union
 
 import numpy as np
 import zarr
@@ -11,10 +12,10 @@ from .validate_patch_dimension import validate_patch_dimensions
 # TOOD split in testable functions
 def extract_patches_random(
     arr: np.ndarray,
-    patch_size: Union[List[int], Tuple[int, ...]],
+    patch_size: Union[list[int], tuple[int, ...]],
     target: Optional[np.ndarray] = None,
     seed: Optional[int] = None,
-) -> Generator[Tuple[np.ndarray, Optional[np.ndarray]], None, None]:
+) -> Generator[tuple[np.ndarray, Optional[np.ndarray]], None, None]:
     """
     Generate patches from an array in a random manner.
 
@@ -31,12 +32,12 @@ def extract_patches_random(
     ----------
     arr : np.ndarray
         Input image array.
-    patch_size : Tuple[int]
+    patch_size : tuple of int
         Patch sizes in each dimension.
     target : Optional[np.ndarray], optional
         Target array, by default None.
-    seed : Optional[int], optional
-        Random seed, by default None.
+    seed : int or None, default=None
+        Random seed.
 
     Yields
     ------
@@ -112,8 +113,8 @@ def extract_patches_random(
 
 def extract_patches_random_from_chunks(
     arr: zarr.Array,
-    patch_size: Union[List[int], Tuple[int, ...]],
-    chunk_size: Union[List[int], Tuple[int, ...]],
+    patch_size: Union[list[int], tuple[int, ...]],
+    chunk_size: Union[list[int], tuple[int, ...]],
     chunk_limit: Optional[int] = None,
     seed: Optional[int] = None,
 ) -> Generator[np.ndarray, None, None]:
@@ -127,9 +128,9 @@ def extract_patches_random_from_chunks(
     ----------
     arr : np.ndarray
         Input image array.
-    patch_size : Union[List[int], Tuple[int, ...]]
+    patch_size : Union[list[int], tuple[int, ...]]
         Patch sizes in each dimension.
-    chunk_size : Union[List[int], Tuple[int, ...]]
+    chunk_size : Union[list[int], tuple[int, ...]]
         Chunk sizes to load from the.
     chunk_limit : Optional[int], optional
         Number of chunks to load, by default None.

careamics/dataset/patching/sequential_patching.py

@@ -1,6 +1,6 @@
 """Sequential patching functions."""
 
-from typing import List, Optional, Tuple, Union
+from typing import Optional, Union
 
 import numpy as np
 from skimage.util import view_as_windows
@@ -9,21 +9,21 @@ from .validate_patch_dimension import validate_patch_dimensions
 
 
 def _compute_number_of_patches(
-    arr_shape: Tuple[int, ...], patch_sizes: Union[List[int], Tuple[int, ...]]
-) -> Tuple[int, ...]:
+    arr_shape: tuple[int, ...], patch_sizes: Union[list[int], tuple[int, ...]]
+) -> tuple[int, ...]:
     """
     Compute the number of patches that fit in each dimension.
 
     Parameters
     ----------
-    arr_shape : Tuple[int, ...]
+    arr_shape : tuple[int, ...]
         Shape of the input array.
-    patch_sizes : Union[List[int], Tuple[int, ...]
+    patch_sizes : Union[list[int], tuple[int, ...]
         Shape of the patches.
 
     Returns
     -------
-    Tuple[int, ...]
+    tuple[int, ...]
         Number of patches in each dimension.
     """
     if len(arr_shape) != len(patch_sizes):
@@ -47,8 +47,8 @@ def _compute_number_of_patches(
 
 
 def _compute_overlap(
-    arr_shape: Tuple[int, ...], patch_sizes: Union[List[int], Tuple[int, ...]]
-) -> Tuple[int, ...]:
+    arr_shape: tuple[int, ...], patch_sizes: Union[list[int], tuple[int, ...]]
+) -> tuple[int, ...]:
     """
     Compute the overlap between patches in each dimension.
 
@@ -57,14 +57,14 @@ def _compute_overlap(
 
     Parameters
     ----------
-    arr_shape : Tuple[int, ...]
+    arr_shape : tuple[int, ...]
         Input array shape.
-    patch_sizes : Union[List[int], Tuple[int, ...]]
+    patch_sizes : Union[list[int], tuple[int, ...]]
         Size of the patches.
 
     Returns
     -------
-    Tuple[int, ...]
+    tuple[int, ...]
         Overlap between patches in each dimension.
     """
     n_patches = _compute_number_of_patches(arr_shape, patch_sizes)
@@ -80,21 +80,21 @@ def _compute_overlap(
 
 
 def _compute_patch_steps(
-    patch_sizes: Union[List[int], Tuple[int, ...]], overlaps: Tuple[int, ...]
-) -> Tuple[int, ...]:
+    patch_sizes: Union[list[int], tuple[int, ...]], overlaps: tuple[int, ...]
+) -> tuple[int, ...]:
     """
     Compute steps between patches.
 
     Parameters
     ----------
-    patch_sizes : Tuple[int]
+    patch_sizes : tuple[int]
         Size of the patches.
-    overlaps : Tuple[int]
+    overlaps : tuple[int]
         Overlap between patches.
 
     Returns
     -------
-    Tuple[int]
+    tuple[int]
         Steps between patches.
     """
     steps = [
@@ -107,9 +107,9 @@ def _compute_patch_steps(
 # TODO why stack the target here and not on a different dimension before this function?
 def _compute_patch_views(
     arr: np.ndarray,
-    window_shape: List[int],
-    step: Tuple[int, ...],
-    output_shape: List[int],
+    window_shape: list[int],
+    step: tuple[int, ...],
+    output_shape: list[int],
     target: Optional[np.ndarray] = None,
 ) -> np.ndarray:
     """
@@ -119,11 +119,11 @@ def _compute_patch_views(
     ----------
     arr : np.ndarray
         Array from which the views are extracted.
-    window_shape : Tuple[int]
+    window_shape : tuple[int]
         Shape of the views.
-    step : Tuple[int]
+    step : tuple[int]
         Steps between views.
-    output_shape : Tuple[int]
+    output_shape : tuple[int]
         Shape of the output array.
     target : Optional[np.ndarray], optional
         Target array, by default None.
@@ -150,9 +150,9 @@ def _compute_patch_views(
 
 def extract_patches_sequential(
     arr: np.ndarray,
-    patch_size: Union[List[int], Tuple[int, ...]],
+    patch_size: Union[list[int], tuple[int, ...]],
     target: Optional[np.ndarray] = None,
-) -> Tuple[np.ndarray, Optional[np.ndarray]]:
+) -> tuple[np.ndarray, Optional[np.ndarray]]:
     """
     Generate patches from an array in a sequential manner.
 
@@ -163,14 +163,14 @@ def extract_patches_sequential(
     ----------
     arr : np.ndarray
         Input image array.
-    patch_size : Tuple[int]
+    patch_size : tuple[int]
         Patch sizes in each dimension.
     target : Optional[np.ndarray], optional
         Target array, by default None.
 
     Returns
     -------
-    Tuple[np.ndarray, Optional[np.ndarray]]
+    tuple[np.ndarray, Optional[np.ndarray]]
         Patches.
     """
     is_3d_patch = len(patch_size) == 3

careamics/dataset/patching/validate_patch_dimension.py

@@ -1,13 +1,13 @@
 """Patch validation functions."""
 
-from typing import List, Tuple, Union
+from typing import Union
 
 import numpy as np
 
 
 def validate_patch_dimensions(
     arr: np.ndarray,
-    patch_size: Union[List[int], Tuple[int, ...]],
+    patch_size: Union[list[int], tuple[int, ...]],
     is_3d_patch: bool,
 ) -> None:
     """
@@ -26,7 +26,7 @@ def validate_patch_dimensions(
     ----------
     arr : np.ndarray
         Input array.
-    patch_size : Union[List[int], Tuple[int, ...]]
+    patch_size : Union[list[int], tuple[int, ...]]
         Size of the patches along each dimension of the array, except the first.
     is_3d_patch : bool
         Whether the patch is 3D or not.

careamics/dataset/tiling/__init__.py

@@ -1,9 +1,9 @@
 """Tiling functions."""
 
 __all__ = [
-    "stitch_prediction",
-    "extract_tiles",
     "collate_tiles",
+    "extract_tiles",
+    "stitch_prediction",
 ]
 
 from .collate_tiles import collate_tiles

careamics/dataset/tiling/collate_tiles.py

@@ -1,6 +1,6 @@
 """Collate function for tiling."""
 
-from typing import Any, List, Tuple
+from typing import Any
 
 import numpy as np
 from torch.utils.data.dataloader import default_collate
@@ -8,7 +8,7 @@ from torch.utils.data.dataloader import default_collate
 from careamics.config.tile_information import TileInformation
 
 
-def collate_tiles(batch: List[Tuple[np.ndarray, TileInformation]]) -> Any:
+def collate_tiles(batch: list[tuple[np.ndarray, TileInformation]]) -> Any:
     """
     Collate tiles received from CAREamics prediction dataloader.
 
@@ -19,7 +19,7 @@ def collate_tiles(batch: List[Tuple[np.ndarray, TileInformation]]) -> Any:
 
     Parameters
     ----------
-    batch : List[Tuple[np.ndarray, TileInformation], ...]
+    batch : list[tuple[np.ndarray, TileInformation], ...]
         Batch of tiles.
 
     Returns

careamics/dataset/tiling/lvae_tiled_patching.py

@@ -2,7 +2,8 @@
 
 import builtins
 import itertools
-from typing import Any, Generator, Optional, Union
+from collections.abc import Generator
+from typing import Any, Optional, Union
 
 import numpy as np
 from numpy.typing import NDArray

careamics/dataset/tiling/tiled_patching.py

@@ -1,7 +1,8 @@
 """Tiled patching utilities."""
 
 import itertools
-from typing import Generator, List, Tuple, Union
+from collections.abc import Generator
+from typing import Union
 
 import numpy as np
 
@@ -10,7 +11,7 @@ from careamics.config.tile_information import TileInformation
 
 def _compute_crop_and_stitch_coords_1d(
     axis_size: int, tile_size: int, overlap: int
-) -> Tuple[List[Tuple[int, int]], List[Tuple[int, int]], List[Tuple[int, int]]]:
+) -> tuple[list[tuple[int, int]], list[tuple[int, int]], list[tuple[int, int]]]:
     """
     Compute the coordinates of each tile along an axis, given the overlap.
 
@@ -25,8 +26,8 @@ def _compute_crop_and_stitch_coords_1d(
 
     Returns
     -------
-    Tuple[Tuple[int, ...], ...]
-        Tuple of all coordinates for given axis.
+    tuple[tuple[int, ...], ...]
+        tuple of all coordinates for given axis.
     """
     # Compute the step between tiles
     step = tile_size - overlap
@@ -81,9 +82,9 @@ def _compute_crop_and_stitch_coords_1d(
 
 def extract_tiles(
     arr: np.ndarray,
-    tile_size: Union[List[int], Tuple[int, ...]],
-    overlaps: Union[List[int], Tuple[int, ...]],
-) -> Generator[Tuple[np.ndarray, TileInformation], None, None]:
+    tile_size: Union[list[int], tuple[int, ...]],
+    overlaps: Union[list[int], tuple[int, ...]],
+) -> Generator[tuple[np.ndarray, TileInformation], None, None]:
     """Generate tiles from the input array with specified overlap.
 
     The tiles cover the whole array. The method returns a generator that yields
@@ -98,14 +99,14 @@ def extract_tiles(
     ----------
     arr : np.ndarray
         Array of shape (S, C, (Z), Y, X).
-    tile_size : Union[List[int], Tuple[int]]
+    tile_size : Union[list[int], tuple[int]]
         Tile sizes in each dimension, of length 2 or 3.
-    overlaps : Union[List[int], Tuple[int]]
+    overlaps : Union[list[int], tuple[int]]
         Overlap values in each dimension, of length 2 or 3.
 
     Yields
     ------
-    Generator[Tuple[np.ndarray, TileInformation], None, None]
+    Generator[tuple[np.ndarray, TileInformation], None, None]
         Tile generator, yields the tile and additional information.
     """
     # Iterate over num samples (S)

careamics/file_io/__init__.py

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

careamics/file_io/read/__init__.py

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

careamics/file_io/read/get_func.py

@@ -1,7 +1,7 @@
 """Module to get read functions."""
 
 from pathlib import Path
-from typing import Callable, Dict, Protocol, Union
+from typing import Callable, Protocol, Union
 
 from numpy.typing import NDArray
 
@@ -30,7 +30,7 @@ class ReadFunc(Protocol):
         """
 
 
-READ_FUNCS: Dict[SupportedData, ReadFunc] = {
+READ_FUNCS: dict[SupportedData, ReadFunc] = {
     SupportedData.TIFF: read_tiff,
 }
 

careamics/file_io/write/__init__.py

@@ -1,10 +1,10 @@
 """Functions relating to writing image files of different formats."""
 
 __all__ = [
+    "SupportedWriteType",
+    "WriteFunc",
     "get_write_func",
     "write_tiff",
-    "WriteFunc",
-    "SupportedWriteType",
 ]
 
 from .get_func import (

careamics/lightning/__init__.py

@@ -2,14 +2,14 @@
 
 __all__ = [
     "FCNModule",
+    "HyperParametersCallback",
+    "PredictDataModule",
+    "ProgressBarCallback",
+    "TrainDataModule",
     "VAEModule",
     "create_careamics_module",
-    "TrainDataModule",
-    "create_train_datamodule",
-    "PredictDataModule",
     "create_predict_datamodule",
-    "HyperParametersCallback",
-    "ProgressBarCallback",
+    "create_train_datamodule",
 ]
 
 from .callbacks import HyperParametersCallback, ProgressBarCallback

careamics/lightning/callbacks/__init__.py

@@ -2,8 +2,8 @@
 
 __all__ = [
     "HyperParametersCallback",
-    "ProgressBarCallback",
     "PredictionWriterCallback",
+    "ProgressBarCallback",
 ]
 
 from .hyperparameters_callback import HyperParametersCallback

careamics/lightning/callbacks/prediction_writer_callback/__init__.py

@@ -1,12 +1,12 @@
 """A package for the `PredictionWriterCallback` class and utilities."""
 
 __all__ = [
+    "CacheTiles",
     "PredictionWriterCallback",
-    "create_write_strategy",
-    "WriteStrategy",
     "WriteImage",
-    "CacheTiles",
+    "WriteStrategy",
     "WriteTilesZarr",
+    "create_write_strategy",
     "select_write_extension",
     "select_write_func",
 ]

careamics/lightning/callbacks/prediction_writer_callback/prediction_writer_callback.py

@@ -2,8 +2,9 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Optional, Sequence, Union
+from typing import Any, Optional, Union
 
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import BasePredictionWriter

careamics/lightning/callbacks/prediction_writer_callback/write_strategy.py

@@ -1,7 +1,8 @@
 """Module containing different strategies for writing predictions."""
 
+from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Optional, Protocol, Sequence, Union
+from typing import Any, Optional, Protocol, Union
 
 import numpy as np
 from numpy.typing import NDArray

careamics/lightning/callbacks/progress_bar_callback.py

@@ -1,7 +1,7 @@
 """Progressbar callback."""
 
 import sys
-from typing import Dict, Union
+from typing import Union
 
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import TQDMProgressBar
@@ -71,7 +71,7 @@ class ProgressBarCallback(TQDMProgressBar):
 
     def get_metrics(
         self, trainer: Trainer, pl_module: LightningModule
-    ) -> Dict[str, Union[int, str, float, Dict[str, float]]]:
+    ) -> dict[str, Union[int, str, float, dict[str, float]]]:
         """Override this to customize the metrics displayed in the progress bar.
 
         Parameters

careamics/lightning/lightning_module.py

@@ -6,7 +6,7 @@ import numpy as np
 import pytorch_lightning as L
 from torch import Tensor, nn
 
-from careamics.config import FCNAlgorithmConfig, VAEAlgorithmConfig
+from careamics.config import UNetBasedAlgorithm, VAEBasedAlgorithm
 from careamics.config.support import (
     SupportedAlgorithm,
     SupportedArchitecture,
@@ -34,6 +34,7 @@ from careamics.utils.torch_utils import get_optimizer, get_scheduler
 NoiseModel = Union[GaussianMixtureNoiseModel, MultiChannelNoiseModel]
 
 
+# TODO rename to UNetModule
 class FCNModule(L.LightningModule):
     """
     CAREamics Lightning module.
@@ -60,7 +61,7 @@ class FCNModule(L.LightningModule):
         Learning rate scheduler name.
     """
 
-    def __init__(self, algorithm_config: Union[FCNAlgorithmConfig, dict]) -> None:
+    def __init__(self, algorithm_config: Union[UNetBasedAlgorithm, dict]) -> None:
         """Lightning module for CAREamics.
 
         This class encapsulates the a PyTorch model along with the training, validation,
@@ -74,7 +75,9 @@ class FCNModule(L.LightningModule):
         super().__init__()
         # if loading from a checkpoint, AlgorithmModel needs to be instantiated
         if isinstance(algorithm_config, dict):
-            algorithm_config = FCNAlgorithmConfig(**algorithm_config)
+            algorithm_config = UNetBasedAlgorithm(
+                **algorithm_config
+            )  # TODO this needs to be updated using the algorithm-specific class
 
         # create model and loss function
         self.model: nn.Module = model_factory(algorithm_config.model)
@@ -266,7 +269,7 @@ class VAEModule(L.LightningModule):
         Learning rate scheduler name.
     """
 
-    def __init__(self, algorithm_config: Union[VAEAlgorithmConfig, dict]) -> None:
+    def __init__(self, algorithm_config: Union[VAEBasedAlgorithm, dict]) -> None:
         """Lightning module for CAREamics.
 
         This class encapsulates the a PyTorch model along with the training, validation,
@@ -280,7 +283,7 @@ class VAEModule(L.LightningModule):
         super().__init__()
         # if loading from a checkpoint, AlgorithmModel needs to be instantiated
         self.algorithm_config = (
-            VAEAlgorithmConfig(**algorithm_config)
+            VAEBasedAlgorithm(**algorithm_config)
             if isinstance(algorithm_config, dict)
             else algorithm_config
         )
@@ -656,9 +659,10 @@ def create_careamics_module(
     algorithm_configuration["model"] = model_configuration
 
     # call the parent init using an AlgorithmModel instance
+    # TODO broken by new configutations!
     algorithm_str = algorithm_configuration["algorithm"]
-    if algorithm_str in FCNAlgorithmConfig.get_compatible_algorithms():
-        return FCNModule(FCNAlgorithmConfig(**algorithm_configuration))
+    if algorithm_str in UNetBasedAlgorithm.get_compatible_algorithms():
+        return FCNModule(UNetBasedAlgorithm(**algorithm_configuration))
     else:
         raise NotImplementedError(
             f"Model {algorithm_str} is not implemented or unknown."