histolab · alessiamarcolini · Nov 15, 2021 · Jul 23, 2021 · Jul 23, 2021 · Jul 26, 2021
diff --git a/histolab/filters/image_filters.py b/histolab/filters/image_filters.py
@@ -765,19 +765,24 @@ class RagThreshold(ImageFilter):
         Color proximity versus space proximity factor. Default is 10.0
     threshold : int, optional
         Threshold value for combining regions. Default is 9.
+    return_labels : bool, optional
+        If True, returns a labeled array where the value denotes segment
+        membership. Otherwise, returns a PIL image where each segment is colored
+        by the average color in it. Default is False.
 
     Returns
     -------
-    PIL.Image.Image
+    PIL.Image.Image, if not ``return_labels``
         Each segment has been colored based on the average
         color for that segment (and similar segments have been combined).
+    np.ndarray, if ``return_labels``
+        Value denotes segment membership.
 
     Raises
     ------
     ValueError
         If ``img`` mode is RGBA.
 
-
     Example:
         >>> from PIL import Image
         >>> from histolab.filters.image_filters import RagThreshold
@@ -787,14 +792,30 @@ class RagThreshold(ImageFilter):
     """  # noqa
 
     def __init__(
-        self, n_segments: int = 800, compactness: float = 10.0, threshold: int = 9
-    ) -> PIL.Image.Image:
+        self,
+        n_segments: int = 800,
+        compactness: float = 10.0,
+        threshold: int = 9,
+        return_labels: bool = False,
+    ) -> None:
         self.n_segments = n_segments
         self.compactness = compactness
         self.threshold = threshold
+        self.return_labels = return_labels
 
-    def __call__(self, img: PIL.Image.Image) -> PIL.Image.Image:
-        return F.rag_threshold(img, self.n_segments, self.compactness, self.threshold)
+    def __call__(
+        self,
+        img: PIL.Image.Image,
+        mask: np.ndarray = None,
+    ) -> Union[PIL.Image.Image, np.ndarray]:
+        return F.rag_threshold(
+            img,
+            n_segments=self.n_segments,
+            compactness=self.compactness,
+            threshold=self.threshold,
+            mask=mask,
+            return_labels=self.return_labels,
+        )
 
 
 class HysteresisThreshold(ImageFilter):

diff --git a/histolab/filters/image_filters_functional.py b/histolab/filters/image_filters_functional.py
@@ -19,7 +19,7 @@
 import math
 import operator
 from functools import reduce
-from typing import Any, Callable
+from typing import Any, Callable, Union
 
 import numpy as np
 import PIL
@@ -413,7 +413,9 @@ def rag_threshold(
     n_segments: int = 800,
     compactness: float = 10.0,
     threshold: int = 9,
-) -> PIL.Image.Image:
+    mask: np.ndarray = None,
+    return_labels: bool = False,
+) -> Union[PIL.Image.Image, np.ndarray]:
     """Combine similar K-means segmented regions based on threshold value.
 
     Segment an image with K-means, build region adjacency graph based on
@@ -430,12 +432,23 @@ def rag_threshold(
         Color proximity versus space proximity factor. Default is 10.0.
     threshold : int, optional
         Threshold value for combining regions. Default is 9.
+    mask : np.ndarray, optional
+        If provided, superpixels are computed only where mask is True,
+        and seed points are homogeneously distributed over the mask
+        using a K-means clustering strategy (See skimage).
+        Must be the same size as ``img``.
+    return_labels : bool, optional
+        If True, returns a labeled array where the value denotes segment
+        membership. Otherwise, returns a PIL image where each segment is colored
+        by the average color in it. Default is False.
 
     Returns
     -------
-    PIL.Image.Image
+    PIL.Image.Image, if not ``return_labels``
         Each segment has been colored based on the average
         color for that segment (and similar segments have been combined).
+    np.ndarray, if ``return_labels``
+        Value denotes segment membership.
 
     Raises
     ------
@@ -445,9 +458,17 @@ def rag_threshold(
     if img.mode == "RGBA":
         raise ValueError("Input image cannot be RGBA")
     img_arr = np.array(img)
-    labels = sk_segmentation.slic(img_arr, n_segments, compactness, start_label=0)
+    labels = sk_segmentation.slic(
+        img_arr,
+        n_segments,
+        compactness,
+        mask=mask,
+        start_label=0 if mask is None else 1,
+    )
     green = sk_future.graph.rag_mean_color(img_arr, labels)
     labels2 = sk_future.graph.cut_threshold(labels, green, threshold)
+    if return_labels:
+        return labels2
     rag = sk_color.label2rgb(labels2, img_arr, kind="avg", bg_label=-1)
     return np_to_pil(rag)
 

diff --git a/tests/expectations/mask-arrays/diagnostic-slide-thumb-rgb-rag-threshold-labels.npy b/tests/expectations/mask-arrays/diagnostic-slide-thumb-rgb-rag-threshold-labels.npy
diff --git a/tests/expectations/mask-arrays/diagnostic-slide-thumb-rgb-rag-threshold-maskedlabels.npy b/tests/expectations/mask-arrays/diagnostic-slide-thumb-rgb-rag-threshold-maskedlabels.npy
diff --git a/tests/integration/test_image_filters.py b/tests/integration/test_image_filters.py
@@ -10,11 +10,19 @@
 from ..util import load_expectation
 
 
+def _create_rag_mask(pil_image):
+    mask = np.ones((pil_image.size[1], pil_image.size[0]))
+    mask[:100, :] = 0
+    mask[:, :100] = 0
+    mask[-100:, :] = 0
+    mask[:, -100:] = 0
+    return mask
+
+
 def test_invert_filter_with_rgb_image():
     expected_value = load_expectation(
         "pil-images-rgb/diagnostic-slide-thumb-rgb-inverted", type_="png"
     )
-
     inverted_img = imf.invert(RGB.DIAGNOSTIC_SLIDE_THUMB_RGB)
 
     np.testing.assert_array_almost_equal(
@@ -484,6 +492,29 @@ def test_rag_threshold_filter_on_rgb_image():
     )
 
 
+@pytest.mark.parametrize(
+    "pil_image, mask, expected_image",
+    (
+        (
+            RGB.DIAGNOSTIC_SLIDE_THUMB_RGB,
+            None,
+            "mask-arrays/diagnostic-slide-thumb-rgb-rag-threshold-labels",
+        ),
+        (
+            RGB.DIAGNOSTIC_SLIDE_THUMB_RGB,
+            _create_rag_mask(RGB.DIAGNOSTIC_SLIDE_THUMB_RGB),
+            "mask-arrays/diagnostic-slide-thumb-rgb-rag-threshold-maskedlabels",
+        ),
+    ),
+)
+def test_rag_threshold_filter_return_labels(pil_image, mask, expected_image):
+    expected_value = load_expectation(expected_image, type_="npy")
+
+    ragged = imf.rag_threshold(pil_image, 650, 20.6, 9, return_labels=True, mask=mask)
+
+    np.testing.assert_array_almost_equal(ragged, expected_value)
+
+
 def test_rag_threshold_filter_on_gs_image():
     gs_img = GS.DIAGNOSTIC_SLIDE_THUMB_GS
     expected_value = load_expectation(

diff --git a/tests/unit/filters/test_image_filters.py b/tests/unit/filters/test_image_filters.py
@@ -162,7 +162,14 @@ def it_calls_rag_threshold_functional(self, request):
 
         rag_threshold(image)
 
-        F_rag_threshold.assert_called_once_with(image, 3, 600, 15)
+        F_rag_threshold.assert_called_once_with(
+            image,
+            n_segments=3,
+            compactness=600,
+            threshold=15,
+            mask=None,
+            return_labels=False,
+        )
         assert type(rag_threshold(image)) == PIL.Image.Image
 
     def it_applies_hysteresis_threshold(self, request):