careamics 0.0.16__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_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}")

careamics/dataset_ng/patch_filter/mask_filter.py

@@ -0,0 +1,95 @@
+"""Filter using an image mask."""
+
+import numpy as np
+
+from careamics.dataset_ng.patch_extractor import GenericImageStack, PatchExtractor
+from careamics.dataset_ng.patch_filter.coordinate_filter_protocol import (
+    CoordinateFilterProtocol,
+)
+from careamics.dataset_ng.patching_strategies import PatchSpecs
+
+
+# TODO is it more intuitive to have a negative mask? (mask of what to avoid)
+class MaskCoordFilter(CoordinateFilterProtocol):
+    """
+    Filter patch coordinates based on an image mask.
+
+    Attributes
+    ----------
+    mask_extractor : PatchExtractor[GenericImageStack]
+        Patch extractor for the binary mask to use for filtering.
+    coverage_perc : float
+        Minimum percentage of masked pixels required to keep a patch.
+    p : float
+        Probability of applying the filter to a patch.
+    rng : np.random.Generator
+        Random number generator for stochastic filtering.
+    """
+
+    def __init__(
+        self,
+        mask_extractor: PatchExtractor[GenericImageStack],
+        coverage: float,
+        p: float = 1.0,
+        seed: int | None = None,
+    ) -> None:
+        """
+        Create a MaskCoordFilter.
+
+        This filter removes patches who fall below a threshold of masked pixels
+        percentage. The mask is expected to be a positive mask where masked pixels
+        correspond to regions of interest.
+
+        Parameters
+        ----------
+        mask_extractor : PatchExtractor[GenericImageStack]
+            The patch extractor for the mask used for filtering.
+        coverage : float
+            Minimum percentage of masked pixels required to keep a patch. Must be
+            between 0 and 1.
+        p : float, default=1
+            Probability of applying the filter to a patch. Must be between 0 and 1.
+        seed : int | None, default=None
+            Seed for the random number generator for reproducibility.
+
+        Raises
+        ------
+        ValueError
+            If coverage is not between 0 and 1.
+        ValueError
+            If p is not between 0 and 1.
+        """
+
+        if not (0 <= coverage <= 1):
+            raise ValueError("Probability p must be between 0 and 1.")
+        if not (0 <= p <= 1):
+            raise ValueError("Probability p must be between 0 and 1.")
+
+        self.mask_extractor = mask_extractor
+        self.coverage = coverage
+
+        self.p = p
+        self.rng = np.random.default_rng(seed)
+
+    def filter_out(self, patch_specs: PatchSpecs) -> bool:
+        """
+        Determine whether to filter out a patch based an image mask.
+
+        Parameters
+        ----------
+        patch : PatchSpecs
+            The patch coordinates to evaluate.
+
+        Returns
+        -------
+        bool
+            True if the patch should be filtered out, False otherwise.
+        """
+
+        if self.rng.uniform(0, 1) < self.p:
+            mask_patch = self.mask_extractor.extract_patch(**patch_specs)
+
+            masked_fraction = np.sum(mask_patch) / mask_patch.size
+            if masked_fraction < self.coverage:
+                return True
+        return False