careamics 0.0.14__py3-none-any.whl → 0.0.16__py3-none-any.whl

careamics/dataset_ng/dataset.py

@@ -1,7 +1,7 @@
 from collections.abc import Sequence
 from enum import Enum
 from pathlib import Path
-from typing import Any, Generic, Literal, NamedTuple, Optional, Union
+from typing import Any, Generic, Literal, NamedTuple, Union
 
 import numpy as np
 from numpy.typing import NDArray
@@ -51,7 +51,7 @@ class CareamicsDataset(Dataset, Generic[GenericImageStack]):
         data_config: NGDataConfig,
         mode: Mode,
         input_extractor: PatchExtractor[GenericImageStack],
-        target_extractor: Optional[PatchExtractor[GenericImageStack]] = None,
+        target_extractor: PatchExtractor[GenericImageStack] | None = None,
     ):
         self.config = data_config
         self.mode = mode
@@ -115,7 +115,7 @@ class CareamicsDataset(Dataset, Generic[GenericImageStack]):
 
         return patching_strategy
 
-    def _initialize_transforms(self) -> Optional[Compose]:
+    def _initialize_transforms(self) -> Compose | None:
         normalize = NormalizeModel(
             image_means=self.input_stats.means,
             image_stds=self.input_stats.stds,

careamics/dataset_ng/factory.py

@@ -1,7 +1,7 @@
 from collections.abc import Sequence
 from enum import Enum
 from pathlib import Path
-from typing import Any, Optional
+from typing import Any
 
 from numpy.typing import NDArray
 from typing_extensions import ParamSpec
@@ -48,8 +48,8 @@ class DatasetType(Enum):
 def determine_dataset_type(
     data_type: SupportedData,
     in_memory: bool,
-    read_func: Optional[ReadFunc] = None,
-    image_stack_loader: Optional[ImageStackLoader] = None,
+    read_func: ReadFunc | None = None,
+    image_stack_loader: ImageStackLoader | None = None,
 ) -> DatasetType:
     """Determine what the dataset type should be based on the input arguments.
 
@@ -121,10 +121,10 @@ def create_dataset(
     inputs: Any,
     targets: Any,
     in_memory: bool,
-    read_func: Optional[ReadFunc] = None,
-    read_kwargs: Optional[dict[str, Any]] = None,
-    image_stack_loader: Optional[ImageStackLoader] = None,
-    image_stack_loader_kwargs: Optional[dict[str, Any]] = None,
+    read_func: ReadFunc | None = None,
+    read_kwargs: dict[str, Any] | None = None,
+    image_stack_loader: ImageStackLoader | None = None,
+    image_stack_loader_kwargs: dict[str, Any] | None = None,
 ) -> CareamicsDataset[ImageStack]:
     """
     Convenience function to create the CAREamicsDataset.
@@ -201,7 +201,7 @@ def create_array_dataset(
     config: NGDataConfig,
     mode: Mode,
     inputs: Sequence[NDArray[Any]],
-    targets: Optional[Sequence[NDArray[Any]]],
+    targets: Sequence[NDArray[Any]] | None,
 ) -> CareamicsDataset[InMemoryImageStack]:
     """
     Create a CAREamicsDataset from array data.
@@ -223,7 +223,7 @@ def create_array_dataset(
         A CAREamicsDataset.
     """
     input_extractor = create_array_extractor(source=inputs, axes=config.axes)
-    target_extractor: Optional[PatchExtractor[InMemoryImageStack]]
+    target_extractor: PatchExtractor[InMemoryImageStack] | None
     if targets is not None:
         target_extractor = create_array_extractor(source=targets, axes=config.axes)
     else:
@@ -235,7 +235,7 @@ def create_tiff_dataset(
     config: NGDataConfig,
     mode: Mode,
     inputs: Sequence[Path],
-    targets: Optional[Sequence[Path]],
+    targets: Sequence[Path] | None,
 ) -> CareamicsDataset[InMemoryImageStack]:
     """
     Create a CAREamicsDataset from tiff files that will be all loaded into memory.
@@ -260,7 +260,7 @@ def create_tiff_dataset(
         source=inputs,
         axes=config.axes,
     )
-    target_extractor: Optional[PatchExtractor[InMemoryImageStack]]
+    target_extractor: PatchExtractor[InMemoryImageStack] | None
     if targets is not None:
         target_extractor = create_tiff_extractor(source=targets, axes=config.axes)
     else:
@@ -273,7 +273,7 @@ def create_czi_dataset(
     config: NGDataConfig,
     mode: Mode,
     inputs: Sequence[Path],
-    targets: Optional[Sequence[Path]],
+    targets: Sequence[Path] | None,
 ) -> CareamicsDataset[CziImageStack]:
     """
     Create a dataset from CZI files.
@@ -296,7 +296,7 @@ def create_czi_dataset(
     """
 
     input_extractor = create_czi_extractor(source=inputs, axes=config.axes)
-    target_extractor: Optional[PatchExtractor[CziImageStack]]
+    target_extractor: PatchExtractor[CziImageStack] | None
     if targets is not None:
         target_extractor = create_czi_extractor(source=targets, axes=config.axes)
     else:
@@ -309,7 +309,7 @@ def create_ome_zarr_dataset(
     config: NGDataConfig,
     mode: Mode,
     inputs: Sequence[Path],
-    targets: Optional[Sequence[Path]],
+    targets: Sequence[Path] | None,
 ) -> CareamicsDataset[ZarrImageStack]:
     """
     Create a dataset from OME ZARR files.
@@ -332,7 +332,7 @@ def create_ome_zarr_dataset(
     """
 
     input_extractor = create_ome_zarr_extractor(source=inputs, axes=config.axes)
-    target_extractor: Optional[PatchExtractor[ZarrImageStack]]
+    target_extractor: PatchExtractor[ZarrImageStack] | None
     if targets is not None:
         target_extractor = create_ome_zarr_extractor(source=targets, axes=config.axes)
     else:
@@ -345,7 +345,7 @@ def create_custom_file_dataset(
     config: NGDataConfig,
     mode: Mode,
     inputs: Sequence[Path],
-    targets: Optional[Sequence[Path]],
+    targets: Sequence[Path] | None,
     *,
     read_func: ReadFunc,
     read_kwargs: dict[str, Any],
@@ -378,7 +378,7 @@ def create_custom_file_dataset(
     input_extractor = create_custom_file_extractor(
         source=inputs, axes=config.axes, read_func=read_func, read_kwargs=read_kwargs
     )
-    target_extractor: Optional[PatchExtractor[InMemoryImageStack]]
+    target_extractor: PatchExtractor[InMemoryImageStack] | None
     if targets is not None:
         target_extractor = create_custom_file_extractor(
             source=targets,
@@ -396,7 +396,7 @@ def create_custom_image_stack_dataset(
     config: NGDataConfig,
     mode: Mode,
     inputs: Any,
-    targets: Optional[Any],
+    targets: Any | None,
     image_stack_loader: ImageStackLoader[P, GenericImageStack],
     *args: P.args,
     **kwargs: P.kwargs,
@@ -436,7 +436,7 @@ def create_custom_image_stack_dataset(
         *args,
         **kwargs,
     )
-    target_extractor: Optional[PatchExtractor[GenericImageStack]]
+    target_extractor: PatchExtractor[GenericImageStack] | None
     if targets is not None:
         target_extractor = create_custom_image_stack_extractor(
             targets,

careamics/dataset_ng/patch_extractor/demo_custom_image_stack_loader.py

@@ -7,7 +7,7 @@ import matplotlib.pyplot as plt
 import numpy as np
 import zarr
 from numpy.typing import NDArray
-from zarr.storage import FSStore
+from zarr.storage import FsspecStore
 
 from careamics.config import DataConfig
 from careamics.config.support import SupportedData
@@ -20,7 +20,7 @@ from careamics.dataset_ng.patch_extractor.patch_extractor_factory import (
 
 # %%
 def create_zarr_array(file_path: Path, data_path: str, data: NDArray):
-    store = FSStore(url=file_path.resolve())
+    store = FsspecStore.from_url(url=file_path.resolve())
     # create array
     array = zarr.create(
         store=store,
@@ -61,7 +61,7 @@ if not file_path.is_file() and not file_path.is_dir():
 # ### Make sure file exists
 
 # %%
-store = FSStore(url=file_path.resolve(), mode="r")
+store = FsspecStore.from_url(url=file_path.resolve(), mode="r")
 
 # %%
 list(store.keys())
@@ -72,7 +72,7 @@ list(store.keys())
 
 # %%
 class ZarrSource(TypedDict):
-    store: FSStore
+    store: FsspecStore
     data_paths: Sequence[str]
 
 

careamics/dataset_ng/patch_extractor/image_stack/in_memory_image_stack.py

@@ -1,9 +1,8 @@
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Literal, Union
+from typing import Any, Literal, Self, Union
 
 from numpy.typing import DTypeLike, NDArray
-from typing_extensions import Self
 
 from careamics.dataset.dataset_utils import reshape_array
 from careamics.file_io.read import ReadFunc, read_tiff

careamics/dataset_ng/patch_extractor/image_stack/zarr_image_stack.py

@@ -1,11 +1,11 @@
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Union
+from typing import Self, Union
 
+import validators
 import zarr
-import zarr.storage
 from numpy.typing import NDArray
-from typing_extensions import Self
+from zarr.storage import FsspecStore, LocalStore
 
 from careamics.dataset.dataset_utils import reshape_array
 
@@ -15,9 +15,10 @@ class ZarrImageStack:
     A class for extracting patches from an image stack that is stored as a zarr array.
     """
 
-    # TODO: keeping store type narrow so that it has the path attribute
-    #   base zarr store is zarr.storage.Store, includes MemoryStore
-    def __init__(self, store: zarr.storage.FSStore, data_path: str, axes: str):
+    # TODO: We should keep store type narrow
+    #   - in zarr v3, does zarr.storage.Store exists and has the path attribute?
+    #   - can we declare a narrow type rather than a union?
+    def __init__(self, store: LocalStore | FsspecStore, data_path: str, axes: str):
         self._store = store
         self._array = zarr.open_array(store=self._store, path=data_path, mode="r")
         # TODO: validate axes
@@ -46,8 +47,33 @@ class ZarrImageStack:
         Assumes the path only contains 1 image.
 
         Path can be to a local file, or it can be a URL to a zarr stored in the cloud.
+
+        Parameters
+        ----------
+        path : Union[Path, str]
+            Path to the root of the OME-Zarr, local file or url.
+
+        Returns
+        -------
+        ZarrImageStack
+            Initialised ZarrImageStack.
+
+        Raises
+        ------
+        ValueError
+            If the path does not exist or is not a valid URL.
+        ValueError
+            If the OME-Zarr at the path does not contain the attribute 'multiscales'.
         """
-        store = zarr.storage.FSStore(url=path)
+        if Path(path).is_file():
+            store = zarr.storage.LocalStore(root=Path(path).resolve())
+        elif validators.url(path):
+            store = zarr.storage.FsspecStore.from_url(url=path)
+        else:
+            raise ValueError(
+                f"Path '{path}' is neither an existing file nor a valid URL."
+            )
+
         group = zarr.open_group(store=store, mode="r")
         if "multiscales" not in group.attrs:
             raise ValueError(

careamics/dataset_ng/patch_extractor/image_stack_loader.py

@@ -38,7 +38,7 @@ class ImageStackLoader(Protocol[P, GenericImageStack]):
 
     >>> from typing import TypedDict
 
-    >>> from zarr.storage import FSStore
+    >>> from zarr.storage import FsspecStore
 
     >>> from careamics.config import DataConfig
     >>> from careamics.dataset_ng.patch_extractor.image_stack import ZarrImageStack
@@ -46,7 +46,7 @@ class ImageStackLoader(Protocol[P, GenericImageStack]):
     >>> # Define a zarr source
     >>> # It encompasses multiple arguments that determine what data will be loaded
     >>> class ZarrSource(TypedDict):
-    ...     store: FSStore
+    ...     store: FsspecStore
     ...     data_paths: Sequence[str]
 
     >>> def custom_image_stack_loader(

careamics/dataset_ng/patching_strategies/random_patching.py

@@ -1,7 +1,6 @@
 """A module for random patching strategies."""
 
 from collections.abc import Sequence
-from typing import Optional
 
 import numpy as np
 
@@ -31,7 +30,7 @@ class RandomPatchingStrategy:
         self,
         data_shapes: Sequence[Sequence[int]],
         patch_size: Sequence[int],
-        seed: Optional[int] = None,
+        seed: int | None = None,
     ):
         """
         A patching strategy for sampling random patches.
@@ -193,7 +192,7 @@ class FixedRandomPatchingStrategy:
         self,
         data_shapes: Sequence[Sequence[int]],
         patch_size: Sequence[int],
-        seed: Optional[int] = None,
+        seed: int | None = None,
     ):
         """A patching strategy for sampling random patches.
 

careamics/dataset_ng/patching_strategies/sequential_patching.py

@@ -1,6 +1,5 @@
 import itertools
 from collections.abc import Sequence
-from typing import Optional
 
 import numpy as np
 from typing_extensions import ParamSpec
@@ -18,7 +17,7 @@ class SequentialPatchingStrategy:
         self,
         data_shapes: Sequence[Sequence[int]],
         patch_size: Sequence[int],
-        overlaps: Optional[Sequence[int]] = None,
+        overlaps: Sequence[int] | None = None,
     ):
         self.data_shapes = data_shapes
         self.patch_size = patch_size

careamics/file_io/read/__init__.py

@@ -9,4 +9,3 @@ __all__ = [
 
 from .get_func import ReadFunc, get_read_func
 from .tiff import read_tiff
-from .zarr import read_zarr

careamics/lightning/__init__.py

@@ -1,18 +1,32 @@
 """CAREamics PyTorch Lightning modules."""
 
 __all__ = [
+    "DataStatsCallback",
     "FCNModule",
     "HyperParametersCallback",
+    "MicroSplitDataModule",
     "PredictDataModule",
     "ProgressBarCallback",
     "TrainDataModule",
     "VAEModule",
     "create_careamics_module",
+    "create_microsplit_predict_datamodule",
+    "create_microsplit_train_datamodule",
     "create_predict_datamodule",
     "create_train_datamodule",
+    "create_unet_based_module",
+    "create_vae_based_module",
 ]
 
-from .callbacks import HyperParametersCallback, ProgressBarCallback
+from .callbacks import DataStatsCallback, HyperParametersCallback, ProgressBarCallback
 from .lightning_module import FCNModule, VAEModule, create_careamics_module
+from .microsplit_data_module import (
+    MicroSplitDataModule,
+    create_microsplit_predict_datamodule,
+    create_microsplit_train_datamodule,
+)
 from .predict_data_module import PredictDataModule, create_predict_datamodule
-from .train_data_module import TrainDataModule, create_train_datamodule
+from .train_data_module import (
+    TrainDataModule,
+    create_train_datamodule,
+)

careamics/lightning/callbacks/__init__.py

@@ -1,11 +1,13 @@
 """Callbacks module."""
 
 __all__ = [
+    "DataStatsCallback",
     "HyperParametersCallback",
     "PredictionWriterCallback",
     "ProgressBarCallback",
 ]
 
+from .data_stats_callback import DataStatsCallback
 from .hyperparameters_callback import HyperParametersCallback
 from .prediction_writer_callback import PredictionWriterCallback
 from .progress_bar_callback import ProgressBarCallback

careamics/lightning/callbacks/data_stats_callback.py

@@ -0,0 +1,23 @@
+"""Data statistics callback."""
+
+import pytorch_lightning as L
+from pytorch_lightning.callbacks import Callback
+
+
+class DataStatsCallback(Callback):
+    """Callback to update model's data statistics from datamodule.
+
+    This callback ensures that the model has access to the data statistics (mean and std)
+    calculated by the datamodule before training starts.
+    """
+
+    def setup(self, trainer: L.Trainer, module: L.LightningModule, stage: str) -> None:
+        """Called when trainer is setting up."""
+        if stage == "fit":
+            # Get data statistics from datamodule
+            (data_mean, data_std), _ = trainer.datamodule.get_data_stats()
+
+            # Set data statistics in the model's likelihood module
+            module.noise_model_likelihood.set_data_stats(
+                data_mean=data_mean["target"], data_std=data_std["target"]
+            )

careamics/lightning/callbacks/prediction_writer_callback/prediction_writer_callback.py

@@ -4,7 +4,7 @@ from __future__ import annotations
 
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Optional, Union
+from typing import Any, Union
 
 from pytorch_lightning import LightningModule, Trainer
 from pytorch_lightning.callbacks import BasePredictionWriter
@@ -84,9 +84,9 @@ class PredictionWriterCallback(BasePredictionWriter):
         cls,
         write_type: SupportedWriteType,
         tiled: bool,
-        write_func: Optional[WriteFunc] = None,
-        write_extension: Optional[str] = None,
-        write_func_kwargs: Optional[dict[str, Any]] = None,
+        write_func: WriteFunc | None = None,
+        write_extension: str | None = None,
+        write_func_kwargs: dict[str, Any] | None = None,
         dirpath: Union[Path, str] = "predictions",
     ) -> PredictionWriterCallback:  # TODO: change type hint to self (find out how)
         """
@@ -172,7 +172,7 @@ class PredictionWriterCallback(BasePredictionWriter):
         trainer: Trainer,
         pl_module: LightningModule,
         prediction: Any,  # TODO: change to expected type
-        batch_indices: Optional[Sequence[int]],
+        batch_indices: Sequence[int] | None,
         batch: Any,  # TODO: change to expected type
         batch_idx: int,
         dataloader_idx: int,

careamics/lightning/callbacks/prediction_writer_callback/write_strategy.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Sequence
 from pathlib import Path
-from typing import Any, Optional, Protocol, Union
+from typing import Any, Protocol, Union
 
 import numpy as np
 from numpy.typing import NDArray
@@ -25,7 +25,7 @@ class WriteStrategy(Protocol):
         trainer: Trainer,
         pl_module: LightningModule,
         prediction: Any,  # TODO: change to expected type
-        batch_indices: Optional[Sequence[int]],
+        batch_indices: Sequence[int] | None,
         batch: Any,  # TODO: change to expected type
         batch_idx: int,
         dataloader_idx: int,
@@ -133,7 +133,7 @@ class CacheTiles(WriteStrategy):
         trainer: Trainer,
         pl_module: LightningModule,
         prediction: tuple[NDArray, list[TileInformation]],
-        batch_indices: Optional[Sequence[int]],
+        batch_indices: Sequence[int] | None,
         batch: tuple[NDArray, list[TileInformation]],
         batch_idx: int,
         dataloader_idx: int,
@@ -259,7 +259,7 @@ class WriteTilesZarr(WriteStrategy):
         trainer: Trainer,
         pl_module: LightningModule,
         prediction: Any,
-        batch_indices: Optional[Sequence[int]],
+        batch_indices: Sequence[int] | None,
         batch: Any,
         batch_idx: int,
         dataloader_idx: int,
@@ -346,7 +346,7 @@ class WriteImage(WriteStrategy):
         trainer: Trainer,
         pl_module: LightningModule,
         prediction: NDArray,
-        batch_indices: Optional[Sequence[int]],
+        batch_indices: Sequence[int] | None,
         batch: NDArray,
         batch_idx: int,
         dataloader_idx: int,

careamics/lightning/callbacks/prediction_writer_callback/write_strategy_factory.py

@@ -1,6 +1,6 @@
 """Module containing convenience function to create `WriteStrategy`."""
 
-from typing import Any, Optional
+from typing import Any
 
 from careamics.config.support import SupportedData
 from careamics.file_io import SupportedWriteType, WriteFunc, get_write_func
@@ -11,9 +11,9 @@ from .write_strategy import CacheTiles, WriteImage, WriteStrategy
 def create_write_strategy(
     write_type: SupportedWriteType,
     tiled: bool,
-    write_func: Optional[WriteFunc] = None,
-    write_extension: Optional[str] = None,
-    write_func_kwargs: Optional[dict[str, Any]] = None,
+    write_func: WriteFunc | None = None,
+    write_extension: str | None = None,
+    write_func_kwargs: dict[str, Any] | None = None,
 ) -> WriteStrategy:
     """
     Create a write strategy from convenient parameters.
@@ -78,8 +78,8 @@ def create_write_strategy(
 
 def _create_tiled_write_strategy(
     write_type: SupportedWriteType,
-    write_func: Optional[WriteFunc],
-    write_extension: Optional[str],
+    write_func: WriteFunc | None,
+    write_extension: str | None,
     write_func_kwargs: dict[str, Any],
 ) -> WriteStrategy:
     """
@@ -130,7 +130,7 @@ def _create_tiled_write_strategy(
 
 
 def select_write_func(
-    write_type: SupportedWriteType, write_func: Optional[WriteFunc] = None
+    write_type: SupportedWriteType, write_func: WriteFunc | None = None
 ) -> WriteFunc:
     """
     Return a function to write images.
@@ -177,7 +177,7 @@ def select_write_func(
 
 
 def select_write_extension(
-    write_type: SupportedWriteType, write_extension: Optional[str] = None
+    write_type: SupportedWriteType, write_extension: str | None = None
 ) -> str:
     """
     Return an extension to add to file paths.