Feature(Next-Gen): Patch filtering rework (#832)

## Disclaimer

- [ ] I am an AI agent.
- [ ] I have used AI and I thoroughly reviewed every line.
- [x] I have not used AI extensively.


## Description

> [!NOTE]  
> **tldr**: This makes the patch filtering in the NG-Dataset more
efficient by using the `StratifiedPatchingStrategy` to set different
sampling regions in the data to have a reduced probability of being
sampled per epoch.

### Background - why do we need this PR?

The old patch filtering implementation used the `RandomPatchingStrategy`
to keep resampling patches until the patch has signal. If there is a
large amount of background in the data this is inefficient. In the new
implementation the filtering process is happens once at the start of
training and there is no need to define a "patience".

### Overview - what changed?

Quite a lot...

- The attributes `patch_filter` and `coord_filter` have been removed
from the `CAREamicsDataset`, and it is also no longer initialised with
masks. This also means all functionality for filtering within the
dataset has been removed.
- Instead, the filtering process happens in `/ng_dataset/filter_bg.py`,
there are two functions, one for filtering with a filter, one for
filtering with a mask.
- The `create_train_val_datasets` and `create_val_split_dataset` factory
functions have been refactored slightly to reduced code duplication. The
`train_dataset` is first created with the new `create_trian_dataset`
within which the data is filtered.
- The `StratifiedPatchingStrategy` has been modified to have a new
method to set the probability that a sampling region will be sampled
from during an epoch. This also necessitates that the number of patches
are reduced if the probability of any regions is reduced.
- `filtered_patch_prob` and `filter_ref_channel` has been added to the
ng data config.
  - The ref channel was added for the case of multiple channels.

### Implementation - how did you implement the changes?

- The idea is a filtering function can be applied to each sampling
region in the `StratifiedPatchingStrategy`, if the region doesn't pass
some threshold the probability of it being sampled from in an epoch can
be reduced.
- The patching strategy instance only needs to store the probability of
each region which doesn't take too much memory.
- When calculating the total probabilities and the sampling bins, the
adjusted areas of the regions can be used, this is the area multiplied
by the new probability.

## Changes Made

### New features or files

- `src/careamics/dataset_ng/filter_bg.py`
  - `filter_background`
  - `filter_background_with_mask`

### Modified features or files

- `src/careamics/dataset_ng/factory.py`
- src/careamics/dataset_ng/patching_strategies/stratified_patching.py

### Removed features or files

- filtering from `CAREamicsDataset`

## How has this been tested?

- Added unit tests that test the number of patches are reduced in the
`StratifiedPatchingStrategy` if the method `set_region_probs` is called.

- Added functional tests for filtering background, from a few different
layers:
  - The patching strategy layer
  - The `filter_background` and `filter_background_with_mask` layer
  - The `create_train_dataset` layer

Still missing is a test for the lightning module layer or an e2e test
with filtering.

## Related Issues

- Resolves #741 

## Breaking changes

- Data config has changed, `filter_patience` has been removed
- Passing masks to the dataset no longer happens

## Additional Notes and Examples

### Still TODO:

- Clean up patch filtering configs and filtering classes, we no longer
need filtering patience or the probability.
- I am considering removing the option to have different "coord filters"
(we don't have any others) because a mask filter config is required to
pass masks to `create_train_dataset`.

---

**Please ensure your PR meets the following requirements:**

- [x] Code builds and passes tests locally, including doctests
- [x] New tests have been added (for bug fixes/features)
- [x] Pre-commit passes
- [ ] PR to the documentation exists (for bug fixes / features)

---------

Co-authored-by: Joran Deschamps <6367888+jdeschamps@users.noreply.github.com>

Author: melisande-c

.pre-commit-config.yaml

@@ -44,7 +44,7 @@ repos:
     rev: v1.10.0
     hooks:
       - id: numpydoc-validation
-        exclude: "^docs/.*|^src/careamics/lvae_training/.*|^src/careamics/models/lvae/.*|^src/careamics/losses/lvae/.*|^scripts/.*|^src/careamics/careamist_v2.py|^src/careamics/lightning/dataset_ng/data_module.py"
+        exclude: "^tests/.*|^docs/.*|^src/careamics/lvae_training/.*|^src/careamics/models/lvae/.*|^src/careamics/losses/lvae/.*|^scripts/.*|^src/careamics/careamist_v2.py|^src/careamics/lightning/dataset_ng/data_module.py"
         # data_module: overloads and config params (val_percentage etc.) don't match signature; see https://github.com/numpy/numpydoc/issues/559
 
   # # jupyter linting and formatting

src/careamics/config/data/ng_data_config.py

@@ -337,9 +337,15 @@ class NGDataConfig(BaseModel):
     """Coordinate filter to apply when using random patching. Only available if
     mode is `training`."""
 
-    patch_filter_patience: int = Field(default=5, ge=1)
-    """Number of consecutive patches not passing the filter before accepting the next
-    patch."""
+    filtered_patch_prob: float = Field(0.1, ge=0.0, le=1.0)
+    """The probability that each patch considered background will be selected each epoch
+    during training. Patches can be considered background by either using a
+    `patch_filter` or by supplying a mask during training. If neither is chosen this
+    parameter is ignored."""
+
+    # TODO: Move inside patch_filter
+    filter_ref_channel: int = 0
+    """The channel to use as reference for filtering."""
 
     augmentations: Sequence[Union[XYFlipConfig, XYRandomRotate90Config]] = Field(
         default=(
@@ -1062,6 +1068,8 @@ def convert_mode(
                     if new_mode == Mode.PREDICTING and new_dataloader_params is not None
                     else self.pred_dataloader_params
                 ),
+                "patch_filter": None,
+                "coord_filter": None,
             }
         )
 

src/careamics/config/data/patch_filter/mask_filter_config.py

@@ -13,5 +13,8 @@ class MaskFilterConfig(FilterConfig):
     name: Literal["mask"] = "mask"
     """Name of the filter."""
 
-    coverage: float = Field(0.5, ge=0.0, le=1.0)
-    """Percentage of masked pixels required to keep a patch."""
+    coverage: float | None = Field(None, ge=0.0, le=1.0)
+    """Minimum ratio of masked pixels required to keep a sampling region,
+    `default=None`. If `None` then `1/(ndims**2)` is used where `ndims` is the number of
+    spatial dimensions.
+    """

src/careamics/dataset_ng/dataset.py

@@ -15,7 +15,6 @@
 from .normalization import create_normalization
 from .normalization.statistics import resolve_normalization_config
 from .patch_extractor import PatchExtractor
-from .patch_filter import create_coord_filter, create_patch_filter
 from .patching_strategies import (
     PatchingStrategy,
     PatchSpecs,
@@ -150,8 +149,6 @@ class CareamicsDataset(Dataset, Generic[GenericImageStack]):
         Extractor for input patches.
     target_extractor : PatchExtractor or None, optional
         Extractor for target patches.
-    mask_extractor : PatchExtractor or None, optional
-        Extractor for mask (e.g. for coord filter).
     """
 
     def __init__(
@@ -160,7 +157,6 @@ def __init__(
         patching_strategy: PatchingStrategy,
         input_extractor: PatchExtractor[GenericImageStack],
         target_extractor: PatchExtractor[GenericImageStack] | None = None,
-        mask_extractor: PatchExtractor[GenericImageStack] | None = None,
     ) -> None:
         """Contructor.
 
@@ -174,8 +170,6 @@ def __init__(
             Extractor for input patches.
         target_extractor : PatchExtractor or None, optional
             Extractor for target patches.
-        mask_extractor : PatchExtractor or None, optional
-            Extractor for mask (e.g. for coord filter).
         """
         # Make sure all the image sizes are greater than the patch size for training
         data_shapes = [
@@ -197,18 +191,6 @@ def __init__(
         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 = patching_strategy
 
         resolve_normalization_config(
@@ -343,54 +325,6 @@ def _extract_patches(
         )
         return input_patch, target_patch
 
-    def _get_filtered_patch(
-        self, index: int
-    ) -> tuple[NDArray[Any], NDArray[Any] | None, PatchSpecs]:
-        """Extract a patch using filtering.
-
-        Parameters
-        ----------
-        index : int
-            Dataset index used to obtain the patch spec.
-
-        Returns
-        -------
-        tuple of (NDArray, NDArray or None, PatchSpecs)
-            Input patch, optional target patch, and patch spec.
-        """
-        should_filter = self.config.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
-
-                    # TODO should we raise an error rather than silently accept patches?
-                    # if patience runs out without ever finding coordinates
-                    # then we need to guard against an exist before defining
-                    # input_patch and target_patch
-                    if patch_filter_patience != 0:
-                        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]]:
@@ -406,7 +340,8 @@ def __getitem__(
         tuple of ImageRegionData
             (input_data,) or (input_data, target_data).
         """
-        input_patch, target_patch, patch_spec = self._get_filtered_patch(index)
+        patch_spec = self.patching_strategy.get_patch_spec(index)
+        input_patch, target_patch = self._extract_patches(patch_spec)
 
         # apply normalization
         input_patch, target_patch = self.normalization(input_patch, target_patch)