careamics 0.0.15__py3-none-any.whl → 0.0.17__py3-none-any.whl

careamics/dataset_ng/dataset.py

@@ -16,6 +16,7 @@ from careamics.config.transformations import NormalizeModel
 from careamics.dataset.dataset_utils.running_stats import WelfordStatistics
 from careamics.dataset.patching.patching import Stats
 from careamics.dataset_ng.patch_extractor import GenericImageStack, PatchExtractor
+from careamics.dataset_ng.patch_filter import create_coord_filter, create_patch_filter
 from careamics.dataset_ng.patching_strategies import (
     FixedRandomPatchingStrategy,
     PatchingStrategy,
@@ -52,13 +53,26 @@ class CareamicsDataset(Dataset, Generic[GenericImageStack]):
         mode: Mode,
         input_extractor: PatchExtractor[GenericImageStack],
         target_extractor: PatchExtractor[GenericImageStack] | None = None,
-    ):
+        mask_extractor: PatchExtractor[GenericImageStack] | None = None,
+    ) -> None:
         self.config = data_config
         self.mode = mode
 
         self.input_extractor = input_extractor
         self.target_extractor = target_extractor
 
+        self.patch_filter = (
+            create_patch_filter(self.config.patch_filter)
+            if self.config.patch_filter is not None
+            else None
+        )
+        self.coord_filter = (
+            create_coord_filter(self.config.coord_filter, mask=mask_extractor)
+            if self.config.coord_filter is not None and mask_extractor is not None
+            else None
+        )
+        self.patch_filter_patience = self.config.patch_filter_patience
+
         self.patching_strategy = self._initialize_patching_strategy()
 
         self.input_stats, self.target_stats = self._initialize_statistics()
@@ -183,10 +197,10 @@ class CareamicsDataset(Dataset, Generic[GenericImageStack]):
             region_spec=patch_spec,
         )
 
-    def __getitem__(
-        self, index: int
-    ) -> Union[tuple[ImageRegionData], tuple[ImageRegionData, ImageRegionData]]:
-        patch_spec = self.patching_strategy.get_patch_spec(index)
+    def _extract_patches(
+        self, patch_spec: PatchSpecs
+    ) -> tuple[NDArray, NDArray | None]:
+        """Extract input and target patches based on patch specifications."""
         input_patch = self.input_extractor.extract_patch(
             data_idx=patch_spec["data_idx"],
             sample_idx=patch_spec["sample_idx"],
@@ -204,7 +218,45 @@ class CareamicsDataset(Dataset, Generic[GenericImageStack]):
             if self.target_extractor is not None
             else None
         )
+        return input_patch, target_patch
+
+    def _get_filtered_patch(
+        self, index: int
+    ) -> tuple[NDArray[Any], NDArray[Any] | None, PatchSpecs]:
+        """Extract a patch that passes filtering criteria with retry logic."""
+        should_filter = self.mode == Mode.TRAINING and (
+            self.patch_filter is not None or self.coord_filter is not None
+        )
+        empty_patch = True
+        patch_filter_patience = self.patch_filter_patience  # reset patience
+
+        while empty_patch and patch_filter_patience > 0:
+            # query patches
+            patch_spec = self.patching_strategy.get_patch_spec(index)
+
+            # filter patch based on coordinates if needed
+            if should_filter and self.coord_filter is not None:
+                if self.coord_filter.filter_out(patch_spec):
+                    patch_filter_patience -= 1
+                    continue
+
+            input_patch, target_patch = self._extract_patches(patch_spec)
+
+            # filter patch based on values if needed
+            if should_filter and self.patch_filter is not None:
+                empty_patch = self.patch_filter.filter_out(input_patch)
+                patch_filter_patience -= 1  # decrease patience
+            else:
+                empty_patch = False
+
+        return input_patch, target_patch, patch_spec
+
+    def __getitem__(
+        self, index: int
+    ) -> Union[tuple[ImageRegionData], tuple[ImageRegionData, ImageRegionData]]:
+        input_patch, target_patch, patch_spec = self._get_filtered_patch(index)
 
+        # apply transforms
         if self.transforms is not None:
             if self.target_extractor is not None:
                 input_patch, target_patch = self.transforms(input_patch, target_patch)

careamics/dataset_ng/factory.py

@@ -121,6 +121,7 @@ def create_dataset(
     inputs: Any,
     targets: Any,
     in_memory: bool,
+    masks: Any = None,
     read_func: ReadFunc | None = None,
     read_kwargs: dict[str, Any] | None = None,
     image_stack_loader: ImageStackLoader | None = None,
@@ -142,6 +143,8 @@ def create_dataset(
     in_memory : bool
         Whether all the data should be loaded into memory. This is argument is ignored
         unless the `data_type` in `config` is "tiff" or "custom".
+    masks : Any, optional
+        The mask sources used to filter patches.
     read_func : ReadFunc, optional
         A function that can that can be used to load custom data. This argument is
         ignored unless the `data_type` in the `config` is "custom".
@@ -168,18 +171,24 @@ def create_dataset(
         data_type, in_memory, read_func, image_stack_loader
     )
     if dataset_type == DatasetType.ARRAY:
-        return create_array_dataset(config, mode, inputs, targets)
+        return create_array_dataset(config, mode, inputs, targets, masks)
     elif dataset_type == DatasetType.IN_MEM_TIFF:
-        return create_tiff_dataset(config, mode, inputs, targets)
+        return create_tiff_dataset(config, mode, inputs, targets, masks)
     # TODO: Lazy tiff
     elif dataset_type == DatasetType.CZI:
-        return create_czi_dataset(config, mode, inputs, targets)
+        return create_czi_dataset(config, mode, inputs, targets, masks)
     elif dataset_type == DatasetType.IN_MEM_CUSTOM_FILE:
         if read_kwargs is None:
             read_kwargs = {}
         assert read_func is not None  # should be true from `determine_dataset_type`
         return create_custom_file_dataset(
-            config, mode, inputs, targets, read_func=read_func, read_kwargs=read_kwargs
+            config,
+            mode,
+            inputs,
+            targets,
+            masks,
+            read_func=read_func,
+            read_kwargs=read_kwargs,
         )
     elif dataset_type == DatasetType.CUSTOM_IMAGE_STACK:
         if image_stack_loader_kwargs is None:
@@ -191,6 +200,7 @@ def create_dataset(
             inputs,
             targets,
             image_stack_loader,
+            masks,
             **image_stack_loader_kwargs,
         )
     else:
@@ -202,6 +212,7 @@ def create_array_dataset(
     mode: Mode,
     inputs: Sequence[NDArray[Any]],
     targets: Sequence[NDArray[Any]] | None,
+    masks: Sequence[NDArray[Any]] | None = None,
 ) -> CareamicsDataset[InMemoryImageStack]:
     """
     Create a CAREamicsDataset from array data.
@@ -216,6 +227,8 @@ def create_array_dataset(
         The input sources to the dataset.
     targets : Any, optional
         The target sources to the dataset.
+    masks : Any, optional
+        The mask sources used to filter patches.
 
     Returns
     -------
@@ -228,7 +241,14 @@ def create_array_dataset(
         target_extractor = create_array_extractor(source=targets, axes=config.axes)
     else:
         target_extractor = None
-    return CareamicsDataset(config, mode, input_extractor, target_extractor)
+    mask_extractor: PatchExtractor[InMemoryImageStack] | None
+    if masks is not None:
+        mask_extractor = create_array_extractor(source=masks, axes=config.axes)
+    else:
+        mask_extractor = None
+    return CareamicsDataset(
+        config, mode, input_extractor, target_extractor, mask_extractor
+    )
 
 
 def create_tiff_dataset(
@@ -236,9 +256,10 @@ def create_tiff_dataset(
     mode: Mode,
     inputs: Sequence[Path],
     targets: Sequence[Path] | None,
+    masks: Sequence[Path] | None = None,
 ) -> CareamicsDataset[InMemoryImageStack]:
     """
-    Create a CAREamicsDataset from tiff files that will be all loaded into memory.
+    Create a CAREamicsDataset from tiff files that will be loaded into memory.
 
     Parameters
     ----------
@@ -246,10 +267,12 @@ def create_tiff_dataset(
         The data configuration.
     mode : Mode
         Whether to create the dataset in "training", "validation" or "predicting" mode.
-    inputs : Any
+    inputs : Sequence[Path]
         The input sources to the dataset.
-    targets : Any, optional
+    targets : Sequence[Path], optional
         The target sources to the dataset.
+    masks : Sequence[Path], optional
+        The mask sources used to filter patches.
 
     Returns
     -------
@@ -265,8 +288,15 @@ def create_tiff_dataset(
         target_extractor = create_tiff_extractor(source=targets, axes=config.axes)
     else:
         target_extractor = None
-    dataset = CareamicsDataset(config, mode, input_extractor, target_extractor)
-    return dataset
+    mask_extractor: PatchExtractor[InMemoryImageStack] | None
+    if masks is not None:
+        mask_extractor = create_tiff_extractor(source=masks, axes=config.axes)
+    else:
+        mask_extractor = None
+
+    return CareamicsDataset(
+        config, mode, input_extractor, target_extractor, mask_extractor
+    )
 
 
 def create_czi_dataset(
@@ -274,6 +304,7 @@ def create_czi_dataset(
     mode: Mode,
     inputs: Sequence[Path],
     targets: Sequence[Path] | None,
+    masks: Sequence[Path] | None = None,
 ) -> CareamicsDataset[CziImageStack]:
     """
     Create a dataset from CZI files.
@@ -288,6 +319,8 @@ def create_czi_dataset(
         The input sources to the dataset.
     targets : Any, optional
         The target sources to the dataset.
+    masks : Any, optional
+        The mask sources used to filter patches.
 
     Returns
     -------
@@ -301,8 +334,15 @@ def create_czi_dataset(
         target_extractor = create_czi_extractor(source=targets, axes=config.axes)
     else:
         target_extractor = None
-    dataset = CareamicsDataset(config, mode, input_extractor, target_extractor)
-    return dataset
+    mask_extractor: PatchExtractor[CziImageStack] | None
+    if masks is not None:
+        mask_extractor = create_czi_extractor(source=masks, axes=config.axes)
+    else:
+        mask_extractor = None
+
+    return CareamicsDataset(
+        config, mode, input_extractor, target_extractor, mask_extractor
+    )
 
 
 def create_ome_zarr_dataset(
@@ -310,6 +350,7 @@ def create_ome_zarr_dataset(
     mode: Mode,
     inputs: Sequence[Path],
     targets: Sequence[Path] | None,
+    masks: Sequence[Path] | None = None,
 ) -> CareamicsDataset[ZarrImageStack]:
     """
     Create a dataset from OME ZARR files.
@@ -324,6 +365,8 @@ def create_ome_zarr_dataset(
         The input sources to the dataset.
     targets : Any, optional
         The target sources to the dataset.
+    masks : Any, optional
+        The mask sources used to filter patches.
 
     Returns
     -------
@@ -337,8 +380,15 @@ def create_ome_zarr_dataset(
         target_extractor = create_ome_zarr_extractor(source=targets, axes=config.axes)
     else:
         target_extractor = None
-    dataset = CareamicsDataset(config, mode, input_extractor, target_extractor)
-    return dataset
+    mask_extractor: PatchExtractor[ZarrImageStack] | None
+    if masks is not None:
+        mask_extractor = create_ome_zarr_extractor(source=masks, axes=config.axes)
+    else:
+        mask_extractor = None
+
+    return CareamicsDataset(
+        config, mode, input_extractor, target_extractor, mask_extractor
+    )
 
 
 def create_custom_file_dataset(
@@ -346,6 +396,7 @@ def create_custom_file_dataset(
     mode: Mode,
     inputs: Sequence[Path],
     targets: Sequence[Path] | None,
+    masks: Sequence[Path] | None = None,
     *,
     read_func: ReadFunc,
     read_kwargs: dict[str, Any],
@@ -363,6 +414,8 @@ def create_custom_file_dataset(
         The input sources to the dataset.
     targets : Any, optional
         The target sources to the dataset.
+    masks : Any, optional
+        The mask sources used to filter patches.
     read_func : Optional[ReadFunc], optional
         A function that can that can be used to load custom data. This argument is
         ignored unless the `data_type` is "custom".
@@ -388,8 +441,21 @@ def create_custom_file_dataset(
         )
     else:
         target_extractor = None
-    dataset = CareamicsDataset(config, mode, input_extractor, target_extractor)
-    return dataset
+
+    mask_extractor: PatchExtractor[InMemoryImageStack] | None
+    if masks is not None:
+        mask_extractor = create_custom_file_extractor(
+            source=masks,
+            axes=config.axes,
+            read_func=read_func,
+            read_kwargs=read_kwargs,
+        )
+    else:
+        mask_extractor = None
+
+    return CareamicsDataset(
+        config, mode, input_extractor, target_extractor, mask_extractor
+    )
 
 
 def create_custom_image_stack_dataset(
@@ -398,6 +464,7 @@ def create_custom_image_stack_dataset(
     inputs: Any,
     targets: Any | None,
     image_stack_loader: ImageStackLoader[P, GenericImageStack],
+    masks: Any | None = None,
     *args: P.args,
     **kwargs: P.kwargs,
 ) -> CareamicsDataset[GenericImageStack]:
@@ -419,6 +486,8 @@ def create_custom_image_stack_dataset(
     image_stack_loader : ImageStackLoader
         A function for custom image stack loading. This argument is ignored unless the
         `data_type` is "custom".
+    masks : Any, optional
+        The mask sources used to filter patches.
     *args : Any
         Positional arguments to pass to the `image_stack_loader`.
     **kwargs : Any
@@ -447,5 +516,19 @@ def create_custom_image_stack_dataset(
         )
     else:
         target_extractor = None
-    dataset = CareamicsDataset(config, mode, input_extractor, target_extractor)
-    return dataset
+
+    mask_extractor: PatchExtractor[GenericImageStack] | None
+    if masks is not None:
+        mask_extractor = create_custom_image_stack_extractor(
+            masks,
+            config.axes,
+            image_stack_loader,
+            *args,
+            **kwargs,
+        )
+    else:
+        mask_extractor = None
+
+    return CareamicsDataset(
+        config, mode, input_extractor, target_extractor, mask_extractor
+    )

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/patch_filter/__init__.py

@@ -0,0 +1,20 @@
+"""Patch filtering strategies."""
+
+__all__ = [
+    "CoordinateFilterProtocol",
+    "MaskCoordFilter",
+    "MaxPatchFilter",
+    "MeanStdPatchFilter",
+    "PatchFilterProtocol",
+    "ShannonPatchFilter",
+    "create_coord_filter",
+    "create_patch_filter",
+]
+
+from .coordinate_filter_protocol import CoordinateFilterProtocol
+from .filter_factory import create_coord_filter, create_patch_filter
+from .mask_filter import MaskCoordFilter
+from .max_filter import MaxPatchFilter
+from .mean_std_filter import MeanStdPatchFilter
+from .patch_filter_protocol import PatchFilterProtocol
+from .shannon_filter import ShannonPatchFilter

careamics/dataset_ng/patch_filter/coordinate_filter_protocol.py

@@ -0,0 +1,27 @@
+"""A protocol for patch filtering."""
+
+from typing import Protocol
+
+from careamics.dataset_ng.patching_strategies import PatchSpecs
+
+
+class CoordinateFilterProtocol(Protocol):
+    """
+    An interface for implementing coordinate filtering strategies.
+    """
+
+    def filter_out(self, patch: PatchSpecs) -> bool:
+        """
+        Determine whether to filter out a given patch based on its coordinates.
+
+        Parameters
+        ----------
+        patch : PatchSpecs
+            The patch coordinates to evaluate.
+
+        Returns
+        -------
+        bool
+            True if the patch should be filtered out (excluded), False otherwise.
+        """
+        ...

careamics/dataset_ng/patch_filter/filter_factory.py

@@ -0,0 +1,94 @@
+"""Factories for coordinate and patch filters."""
+
+from typing import Union
+
+from careamics.config.data.patch_filter import (
+    FilterModel,
+    MaskFilterModel,
+    MaxFilterModel,
+    MeanSTDFilterModel,
+    ShannonFilterModel,
+)
+from careamics.config.support.supported_filters import (
+    SupportedCoordinateFilters,
+    SupportedPatchFilters,
+)
+from careamics.dataset_ng.patch_extractor import GenericImageStack, PatchExtractor
+
+from .mask_filter import MaskCoordFilter
+from .max_filter import MaxPatchFilter
+from .mean_std_filter import MeanStdPatchFilter
+from .shannon_filter import ShannonPatchFilter
+
+PatchFilter = Union[
+    MaxPatchFilter,
+    MeanStdPatchFilter,
+    ShannonPatchFilter,
+]
+
+
+CoordFilter = Union[MaskCoordFilter]
+
+
+def create_coord_filter(
+    filter_model: FilterModel, mask: PatchExtractor[GenericImageStack]
+) -> CoordFilter:
+    """Factory function to create coordinate filter instances based on the filter name.
+
+    Parameters
+    ----------
+    filter_model : FilterModel
+        Pydantic model of the filter to be created.
+    mask : PatchExtractor[GenericImageStack]
+        Mask extractor to be used for the mask filter.
+
+    Returns
+    -------
+    CoordFilter
+        Instance of the mask patch filter.
+    """
+    if filter_model.name == SupportedCoordinateFilters.MASK:
+        assert isinstance(filter_model, MaskFilterModel)
+        return MaskCoordFilter(
+            mask_extractor=mask,
+            coverage=filter_model.coverage,
+            p=filter_model.p,
+            seed=filter_model.seed,
+        )
+    else:
+        raise ValueError(f"Unknown filter name: {filter_model}")
+
+
+def create_patch_filter(filter_model: FilterModel) -> PatchFilter:
+    """Factory function to create patch filter instances based on the filter name.
+
+    Parameters
+    ----------
+    filter_model : FilterModel
+        Pydantic model of the filter to be created.
+
+    Returns
+    -------
+    PatchFilter
+        Instance of the requested patch filter.
+    """
+    if filter_model.name == SupportedPatchFilters.MAX:
+        assert isinstance(filter_model, MaxFilterModel)
+        return MaxPatchFilter(
+            threshold=filter_model.threshold, p=filter_model.p, seed=filter_model.seed
+        )
+    elif filter_model.name == SupportedPatchFilters.MEANSTD:
+        assert isinstance(filter_model, MeanSTDFilterModel)
+        return MeanStdPatchFilter(
+            mean_threshold=filter_model.mean_threshold,
+            std_threshold=filter_model.std_threshold,
+            p=filter_model.p,
+            seed=filter_model.seed,
+        )
+    elif filter_model.name == SupportedPatchFilters.SHANNON:
+        assert isinstance(filter_model, ShannonFilterModel)
+        return ShannonPatchFilter(
+            threshold=filter_model.threshold, p=filter_model.p, seed=filter_model.seed
+        )
+    else:
+        raise ValueError(f"Unknown filter name: {filter_model}")