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

careamics/dataset_ng/patch_filter/shannon_filter.py

@@ -0,0 +1,188 @@
+"""Filter patches based on Shannon entropy threshold."""
+
+from collections.abc import Sequence
+
+import numpy as np
+from skimage.measure import shannon_entropy
+from tqdm import tqdm
+
+from careamics.dataset_ng.patch_extractor.patch_extractor_factory import (
+    create_array_extractor,
+)
+from careamics.dataset_ng.patch_filter.patch_filter_protocol import PatchFilterProtocol
+from careamics.dataset_ng.patching_strategies import TilingStrategy
+
+
+class ShannonPatchFilter(PatchFilterProtocol):
+    """
+    Filter patches based on Shannon entropy threshold.
+
+    Attributes
+    ----------
+    threshold : float
+        Threshold for the Shannon entropy of the patch.
+    p : float
+        Probability of applying the filter to a patch.
+    rng : np.random.Generator
+        Random number generator for stochastic filtering.
+    """
+
+    def __init__(
+        self, threshold: float, p: float = 1.0, seed: int | None = None
+    ) -> None:
+        """
+        Create a ShannonEntropyFilter.
+
+        This filter removes patches whose Shannon entropy is below a specified
+        threshold.
+
+        Parameters
+        ----------
+        threshold : float
+            Threshold for the Shannon entropy of the patch.
+        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 threshold is negative.
+        ValueError
+            If p is not between 0 and 1.
+        """
+        if threshold < 0:
+            raise ValueError("Threshold must be non-negative.")
+        if not (0 <= p <= 1):
+            raise ValueError("Probability p must be between 0 and 1.")
+
+        self.threshold = threshold
+
+        self.p = p
+        self.rng = np.random.default_rng(seed)
+
+    def filter_out(self, patch: np.ndarray) -> bool:
+        """
+        Determine whether to filter out a patch based on its Shannon entropy.
+
+        Parameters
+        ----------
+        patch : numpy.NDArray
+            The patch to evaluate.
+
+        Returns
+        -------
+        bool
+            True if the patch should be filtered out, False otherwise.
+        """
+        if self.rng.uniform(0, 1) < self.p:
+            return shannon_entropy(patch) < self.threshold
+        return False
+
+    @staticmethod
+    def filter_map(
+        image: np.ndarray,
+        patch_size: Sequence[int],
+    ) -> np.ndarray:
+        """
+        Compute the Shannon entropy map of an image.
+
+        The entropy is computed over non-overlapping patches. This method can be used
+        to assess a useful threshold for the Shannon entropy filter.
+
+        Parameters
+        ----------
+        image : numpy.NDArray
+            The image for which to compute the entropy map, must be 2D or 3D.
+        patch_size : Sequence[int]
+            The size of the patches to compute the entropy over. Must be a sequence
+            of two integers.
+
+        Returns
+        -------
+        numpy.NDArray
+            The Shannon entropy map of the patch.
+
+        Raises
+        ------
+        ValueError
+            If the image is not 2D or 3D.
+
+        Example
+        -------
+        The `filter_map` method can be used to assess a useful threshold for the
+        Shannon entropy filter. Below is an example of how to compute and visualize
+        the Shannon entropy map of a random image and visualize thresholded versions
+        of the map.
+        >>> import numpy as np
+        >>> from matplotlib import pyplot as plt
+        >>> from careamics.dataset_ng.patch_filter import ShannonPatchFilter
+        >>> rng = np.random.default_rng(42)
+        >>> image = rng.binomial(20, 0.1, (256, 256)).astype(np.float32)
+        >>> image[64:192, 64:192] += rng.normal(50, 5, (128, 128))
+        >>> image[96:160, 96:160] = rng.poisson(image[96:160, 96:160])
+        >>> patch_size = (16, 16)
+        >>> entropy_map = ShannonPatchFilter.filter_map(image, patch_size)
+        >>> fig, ax = plt.subplots(1, 5, figsize=(20, 5)) # doctest: +SKIP
+        >>> for i, thresh in enumerate([2 + 1.5 * i for i in range(5)]):
+        ...     ax[i].imshow(entropy_map >= thresh, cmap="gray") #doctest: +SKIP
+        ...     ax[i].set_title(f"Threshold: {thresh}") #doctest: +SKIP
+        >>> plt.show() # doctest: +SKIP
+        """
+        if len(image.shape) < 2 or len(image.shape) > 3:
+            raise ValueError("Image must be 2D or 3D.")
+
+        axes = "YX" if len(patch_size) == 2 else "ZYX"
+
+        shannon_img = np.zeros_like(image, dtype=float)
+
+        extractor = create_array_extractor(source=[image], axes=axes)
+        tiling = TilingStrategy(
+            data_shapes=[(1, 1, *image.shape)],
+            tile_size=patch_size,
+            overlaps=(0,) * len(patch_size),  # no overlap
+        )
+
+        for idx in tqdm(range(tiling.n_patches), desc="Computing Shannon Entropy map"):
+            patch_spec = tiling.get_patch_spec(idx)
+            patch = extractor.extract_patch(
+                data_idx=0,
+                sample_idx=0,
+                coords=patch_spec["coords"],
+                patch_size=patch_size,
+            )
+
+            coordinates = tuple(
+                slice(patch_spec["coords"][i], patch_spec["coords"][i] + p)
+                for i, p in enumerate(patch_size)
+            )
+            shannon_img[coordinates] = shannon_entropy(patch)
+
+        return shannon_img
+
+    @staticmethod
+    def apply_filter(
+        filter_map: np.ndarray,
+        threshold: float,
+    ) -> np.ndarray:
+        """
+        Apply the Shannon entropy filter to a precomputed filter map.
+
+        The filter map is the output of the `filter_map` method.
+
+        Parameters
+        ----------
+        filter_map : numpy.NDArray
+            The precomputed Shannon entropy map of the image.
+        threshold : float
+            The Shannon entropy threshold for filtering.
+
+        Returns
+        -------
+        numpy.NDArray
+            A boolean array where True indicates that the patch should be kept
+            (not filtered out) and False indicates that the patch should be filtered
+            out.
+        """
+        return filter_map >= threshold

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,33 @@
+"""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.
+
+        Parameters
+        ----------
+        trainer : Lightning.Trainer
+            The trainer instance.
+        module : Lightning.LightningModule
+            The model being trained.
+        stage : str
+            The current stage of training (e.g., 'fit', 'validate', 'test', 'predict').
+        """
+        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/dataset_ng/data_module.py

@@ -39,6 +39,10 @@ class CareamicsDataModule(L.LightningDataModule):
     train_data_target : Optional[InputType]
         Training data target, can be a path to a folder,
         a list of paths, or a numpy array.
+    train_data_mask : InputType (when filtering is needed)
+        Training data mask, can be a path to a folder,
+        a list of paths, or a numpy array. Used for coordinate filtering.
+        Only required when using coordinate-based patch filtering.
     val_data : Optional[InputType]
         Validation data, can be a path to a folder,
         a list of paths, or a numpy array.
@@ -99,6 +103,9 @@ class CareamicsDataModule(L.LightningDataModule):
     train_data_target : Optional[Any]
         Training data target, can be a path to a folder, a list of paths, or a numpy
         array.
+    train_data_mask : Optional[Any]
+        Training data mask, can be a path to a folder, a list of paths, or a numpy
+        array.
     val_data : Optional[Any]
         Validation data, can be a path to a folder, a list of paths, or a numpy array.
     val_data_target : Optional[Any]
@@ -118,7 +125,7 @@ class CareamicsDataModule(L.LightningDataModule):
         If input and target data types are not consistent.
     """
 
-    # standard use
+    # standard use (no mask)
     @overload
     def __init__(
         self,
@@ -136,7 +143,26 @@ class CareamicsDataModule(L.LightningDataModule):
         use_in_memory: bool = True,
     ) -> None: ...
 
-    # custom read function
+    # with training mask for filtering
+    @overload
+    def __init__(
+        self,
+        data_config: NGDataConfig,
+        *,
+        train_data: InputType | None = None,
+        train_data_target: InputType | None = None,
+        train_data_mask: InputType,
+        val_data: InputType | None = None,
+        val_data_target: InputType | None = None,
+        pred_data: InputType | None = None,
+        pred_data_target: InputType | None = None,
+        extension_filter: str = "",
+        val_percentage: float | None = None,
+        val_minimum_split: int = 5,
+        use_in_memory: bool = True,
+    ) -> None: ...
+
+    # custom read function (no mask)
     @overload
     def __init__(
         self,
@@ -156,6 +182,48 @@ class CareamicsDataModule(L.LightningDataModule):
         use_in_memory: bool = True,
     ) -> None: ...
 
+    # custom read function with training mask
+    @overload
+    def __init__(
+        self,
+        data_config: NGDataConfig,
+        *,
+        train_data: InputType | None = None,
+        train_data_target: InputType | None = None,
+        train_data_mask: InputType,
+        val_data: InputType | None = None,
+        val_data_target: InputType | None = None,
+        pred_data: InputType | None = None,
+        pred_data_target: InputType | None = None,
+        read_source_func: Callable,
+        read_kwargs: dict[str, Any] | None = None,
+        extension_filter: str = "",
+        val_percentage: float | None = None,
+        val_minimum_split: int = 5,
+        use_in_memory: bool = True,
+    ) -> None: ...
+
+    # image stack loader (no mask)
+    @overload
+    def __init__(
+        self,
+        data_config: NGDataConfig,
+        *,
+        train_data: Any | None = None,
+        train_data_target: Any | None = None,
+        val_data: Any | None = None,
+        val_data_target: Any | None = None,
+        pred_data: Any | None = None,
+        pred_data_target: Any | None = None,
+        image_stack_loader: ImageStackLoader,
+        image_stack_loader_kwargs: dict[str, Any] | None = None,
+        extension_filter: str = "",
+        val_percentage: float | None = None,
+        val_minimum_split: int = 5,
+        use_in_memory: bool = True,
+    ) -> None: ...
+
+    # image stack loader with training mask
     @overload
     def __init__(
         self,
@@ -163,6 +231,7 @@ class CareamicsDataModule(L.LightningDataModule):
         *,
         train_data: Any | None = None,
         train_data_target: Any | None = None,
+        train_data_mask: Any,
         val_data: Any | None = None,
         val_data_target: Any | None = None,
         pred_data: Any | None = None,
@@ -181,6 +250,7 @@ class CareamicsDataModule(L.LightningDataModule):
         *,
         train_data: Any | None = None,
         train_data_target: Any | None = None,
+        train_data_mask: Any | None = None,
         val_data: Any | None = None,
         val_data_target: Any | None = None,
         pred_data: Any | None = None,
@@ -209,6 +279,10 @@ class CareamicsDataModule(L.LightningDataModule):
         train_data_target : Optional[InputType]
             Training data target, can be a path to a folder,
             a list of paths, or a numpy array.
+        train_data_mask : InputType (when filtering is needed)
+            Training data mask, can be a path to a folder,
+            a list of paths, or a numpy array. Used for coordinate filtering.
+            Only required when using coordinate-based patch filtering.
         val_data : Optional[InputType]
             Validation data, can be a path to a folder,
             a list of paths, or a numpy array.
@@ -268,6 +342,8 @@ class CareamicsDataModule(L.LightningDataModule):
         self.train_data, self.train_data_target = self._initialize_data_pair(
             train_data, train_data_target
         )
+        self.train_data_mask, _ = self._initialize_data_pair(train_data_mask, None)
+
         self.val_data, self.val_data_target = self._initialize_data_pair(
             val_data, val_data_target
         )
@@ -574,6 +650,7 @@ class CareamicsDataModule(L.LightningDataModule):
                 mode=Mode.TRAINING,
                 inputs=self.train_data,
                 targets=self.train_data_target,
+                masks=self.train_data_mask,
                 config=self.config,
                 in_memory=self.use_in_memory,
                 read_func=self.read_source_func,