🆕 Use Annotations as a Filter for Patch Extraction

tests/test_patch_extraction.py

@@ -7,7 +7,9 @@
 
 import numpy as np
 import pytest
+from shapely.geometry import Polygon
 
+from tiatoolbox.annotation.storage import Annotation, SQLiteStore
 from tiatoolbox.tools import patchextraction
 from tiatoolbox.tools.patchextraction import PatchExtractor
 from tiatoolbox.utils import misc
@@ -322,7 +324,7 @@ def test_get_coordinates() -> None:
     )
     # test when output patch shape is out of bound
     # but input is in bound
-    input_bounds, output_bounds = PatchExtractor.get_coordinates(
+    input_bounds, output_bounds = PatchExtractor.get_coordinates(  # skipcq: PYL-E0633
         image_shape=(9, 6),
         patch_input_shape=(5, 5),
         patch_output_shape=(4, 4),
@@ -512,6 +514,7 @@ def test_filter_coordinates() -> None:
 def test_mask_based_patch_extractor_ndpi(
     sample_ndpi: Path,
     caplog: pytest.LogCaptureFixture,
+    tmp_path: Path,
 ) -> None:
     """Test SlidingWindowPatchExtractor with mask for ndpi image."""
     res = 0
@@ -607,3 +610,56 @@ def test_mask_based_patch_extractor_ndpi(
         stride=stride,
     )
     assert "No candidate coordinates left" in caplog.text
+
+    # test passing an annotation mask
+    ann = Annotation(
+        Polygon.from_bounds(0, 0, slide_dimensions[0], int(slide_dimensions[1] / 4)),
+        {"label": "region1"},
+    )
+    ann2 = Annotation(
+        Polygon.from_bounds(
+            0, int(slide_dimensions[1] / 2), slide_dimensions[0], slide_dimensions[1]
+        ),
+        {"label": "region2"},
+    )
+    store = SQLiteStore(tmp_path / "test.db")
+    store.append_many([ann, ann2])
+    store.close()
+
+    patches = patchextraction.get_patch_extractor(
+        input_img=input_img,
+        input_mask=str(tmp_path / "test.db"),
+        method_name="slidingwindow",
+        patch_size=patch_size,
+        resolution=res,
+        units="level",
+        stride=None,
+        store_filter=None,
+    )
+    len_all = len(patches)
+
+    patches = patchextraction.get_patch_extractor(
+        input_img=input_img,
+        input_mask=str(tmp_path / "test.db"),
+        method_name="slidingwindow",
+        patch_size=patch_size,
+        resolution=res,
+        units="level",
+        stride=None,
+        store_filter="props['label'] == 'region2'",
+    )
+    len_region2 = len(patches)
+
+    patches = patchextraction.get_patch_extractor(
+        input_img=input_img,
+        input_mask=str(tmp_path / "test.db"),
+        method_name="slidingwindow",
+        patch_size=patch_size,
+        resolution=res,
+        units="level",
+        stride=None,
+        store_filter="props['label'] == 'region1'",
+    )
+    len_region1 = len(patches)
+
+    assert len_all > len_region2 > len_region1

tiatoolbox/data/remote_samples.yaml

@@ -143,5 +143,7 @@ files:
     url: [ *testdata, "annotation/test1_config.json"]
   config_2:
     url: [ *testdata, "annotation/test2_config.json"]
+  patch_annotations:
+    url: [ *testdata, "annotation/sample_wsi_patch_preds.db"]
   nuclick-output:
     url: [*modelroot, "predictions/nuclei_mask/nuclick-output.npy"]

tiatoolbox/tools/patchextraction.py

@@ -11,13 +11,15 @@
 from tiatoolbox import logger
 from tiatoolbox.utils import misc
 from tiatoolbox.utils.exceptions import MethodNotSupportedError
+from tiatoolbox.utils.visualization import AnnotationRenderer
 from tiatoolbox.wsicore import wsireader
 
 if TYPE_CHECKING:  # pragma: no cover
     from pathlib import Path
 
     from pandas import DataFrame
 
+    from tiatoolbox.annotation.storage import AnnotationStore
     from tiatoolbox.typing import Resolution, Units
 
 
@@ -45,9 +47,12 @@ class ExtractorParams(TypedDict, total=False):
     pad_mode: str
     pad_constant_values: int | tuple[int, int]
     within_bound: bool
-    input_mask: str | Path | np.ndarray | wsireader.VirtualWSIReader
+    input_mask: (
+        str | Path | np.ndarray | wsireader.VirtualWSIReader | AnnotationStore | None
+    )
     stride: int | tuple[int, int]
     min_mask_ratio: float
+    store_filter: str | None
 
 
 class PointsPatchExtractorParams(TypedDict):
@@ -81,9 +86,12 @@ class SlidingWindowPatchExtractorParams(TypedDict):
     pad_mode: str
     pad_constant_values: int | tuple[int, int]
     within_bound: bool
-    input_mask: str | Path | np.ndarray | wsireader.VirtualWSIReader | None
+    input_mask: (
+        str | Path | np.ndarray | wsireader.VirtualWSIReader | AnnotationStore | None
+    )
     stride: int | tuple[int, int] | None
     min_mask_ratio: float
+    store_filter: str | None
 
 
 class PatchExtractorABC(ABC):
@@ -123,7 +131,10 @@ class PatchExtractor(PatchExtractorABC):
             'morphological' options. In case of 'otsu' or
             'morphological', a tissue mask is generated for the
             input_image using tiatoolbox :class:`TissueMasker`
-            functionality.
+            functionality. May also be an annotation store, in which case the
+            mask is generated based on the annotations. All annotations are used by
+            default; the 'store_filter' argument can be used to specify a filter for
+            a subset of annotations to use to build the mask.
         resolution (Resolution):
             Resolution at which to read the image, default = 0. Either a
             single number or a sequence of two numbers for x and y are
@@ -150,6 +161,10 @@ class PatchExtractor(PatchExtractorABC):
         min_mask_ratio (float):
             Area in percentage that a patch needs to contain of positive
             mask to be included. Defaults to 0.
+        store_filter (str):
+            Filter to apply to the annotations when generating the mask. Default is
+            None, which uses all annotations. Only used if the provided mask is an
+            annotation store.
 
 
     Attributes:
@@ -188,12 +203,18 @@ def __init__(
         self: PatchExtractor,
         input_img: str | Path | np.ndarray,
         patch_size: int | tuple[int, int],
-        input_mask: str | Path | np.ndarray | wsireader.VirtualWSIReader | None = None,
+        input_mask: str
+        | Path
+        | np.ndarray
+        | wsireader.VirtualWSIReader
+        | AnnotationStore
+        | None = None,
         resolution: Resolution = 0,
         units: Units = "level",
         pad_mode: str = "constant",
         pad_constant_values: int | tuple[int, int] = 0,
         min_mask_ratio: float = 0,
+        store_filter: str | None = None,
         *,
         within_bound: bool = False,
     ) -> None:
@@ -216,6 +237,22 @@ def __init__(
 
         if input_mask is None:
             self.mask = None
+        elif isinstance(input_mask, str) and input_mask.endswith(".db"):
+            # input_mask is an annotation store
+            renderer = AnnotationRenderer(
+                max_scale=10000, edge_thickness=0, where=store_filter
+            )
+            rendered_mask = wsireader.AnnotationStoreReader(
+                input_mask,
+                renderer=renderer,
+                info=self.wsi.info,
+            ).slide_thumbnail()
+            rendered_mask = rendered_mask[:, :, 0] == 0
+            self.mask = wsireader.VirtualWSIReader(
+                rendered_mask,
+                info=self.wsi.info,
+                mode="bool",
+            )
         elif isinstance(input_mask, str) and input_mask in {"otsu", "morphological"}:
             if isinstance(self.wsi, wsireader.VirtualWSIReader):
                 self.mask = None
@@ -618,14 +655,18 @@ class SlidingWindowPatchExtractor(PatchExtractor):
         min_mask_ratio (float):
             Only patches with positive area percentage above this value are included.
             Defaults to 0.
+        store_filter (str):
+            Filter to apply to the annotations when generating the mask. Default is
+            None, which uses all annotations. Only used if the provided mask is an
+            annotation store.
 
     Attributes:
         stride(tuple(int)):
             Stride in (x, y) direction for patch extraction.
 
     """
 
-    def __init__(
+    def __init__(  # noqa: PLR0913
         self: SlidingWindowPatchExtractor,
         input_img: str | Path | np.ndarray,
         patch_size: int | tuple[int, int],
@@ -636,6 +677,7 @@ def __init__(
         pad_mode: str = "constant",
         pad_constant_values: int | tuple[int, int] = 0,
         min_mask_ratio: float = 0,
+        store_filter: str | None = None,
         *,
         within_bound: bool = False,
     ) -> None:
@@ -650,6 +692,7 @@ def __init__(
             pad_constant_values=pad_constant_values,
             within_bound=within_bound,
             min_mask_ratio=min_mask_ratio,
+            store_filter=store_filter,
         )
         if stride is None:
             self.stride = self.patch_size
@@ -794,5 +837,6 @@ def get_patch_extractor(
         "pad_constant_values": kwargs.get("pad_constant_values", 0),
         "min_mask_ratio": kwargs.get("min_mask_ratio", 0),
         "within_bound": kwargs.get("within_bound", False),
+        "store_filter": kwargs.get("store_filter"),
     }
     return SlidingWindowPatchExtractor(**sliding_window_patch_extractor_args)