Merge PR #202 | A generic pmap-guided object removal widget

Author: qin-yu

plantseg/dataprocessing/functional/advanced_dataprocessing.py

@@ -5,8 +5,8 @@
 import numpy as np
 import tqdm
 from skimage.filters import gaussian
-from skimage.segmentation import watershed
-
+from skimage.segmentation import watershed, relabel_sequential
+from skimage.measure import regionprops
 
 def get_bbox(mask: np.array, pixel_toll: int = 0) -> tuple[tuple, int, int, int]:
     """
@@ -235,3 +235,52 @@ def fix_over_under_segmentation_from_nuclei(cell_seg: np.array,
                                        cell_assignments,
                                        cell_idx=None)
     return _cell_seg
+
+
+def remove_false_positives_by_foreground_probability(segmentation: np.array,
+                                                     foreground: np.array,
+                                                     threshold: float) -> np.array:
+    """
+    Remove false positive regions in a segmentation based on a foreground probability map in a smart way.
+    If the mean(an instance * its own probability region) < threshold, it is removed.
+
+    Args:
+        segmentation (np.ndarray): The segmentation array, where each unique non-zero value indicates a distinct region.
+        foreground (np.ndarray): The foreground probability map, same shape as `segmentation`.
+        threshold (float): Probability threshold below which regions are considered false positives.
+
+    Returns:
+        np.ndarray: The modified segmentation array with false positives removed.
+    """
+    # TODO: make a channel for removed regions for easier inspection
+    # TODO: use `relabel_sequential` to recover the original labels
+
+    if not segmentation.shape == foreground.shape:
+        raise ValueError("Shape of segmentation and probability map must match.")
+    if foreground.max() > 1:
+        raise ValueError("Foreground must be a probability map probability map.")
+
+    instances, _, _ = relabel_sequential(segmentation)
+
+    regions = regionprops(instances)
+    to_keep = np.ones(len(regions) + 1)
+    pixel_count = np.zeros(len(regions) + 1)
+    pixel_value = np.zeros(len(regions) + 1)
+
+    for region in tqdm.tqdm(regions):
+        bbox = region.bbox
+        cube = instances[bbox[0]:bbox[3], bbox[1]:bbox[4], bbox[2]:bbox[5]] == region.label  # other instances may exist, don't use `> 0`
+        prob = foreground[bbox[0]:bbox[3], bbox[1]:bbox[4], bbox[2]:bbox[5]]
+        pixel_count[region.label] = region.area
+        pixel_value[region.label] = (cube * prob).sum()
+
+    likelihood = pixel_value / pixel_count
+    to_keep[likelihood < threshold] = 0
+    ids_to_delete = np.argwhere(to_keep == 0)
+    assert ids_to_delete.shape[1] == 1
+    ids_to_delete = ids_to_delete.flatten()
+    # print(f"    Removing instance {region.label}: pixel count: {pixel_count}, pixel value: {pixel_value}, likelihood: {likelihood}")
+
+    instances[np.isin(instances, ids_to_delete)] = 0
+    instances, _, _ = relabel_sequential(instances)
+    return instances

plantseg/segmentation/gasp.py

@@ -56,9 +56,7 @@ def __init__(self,
 
         # Postprocessing size threshold
         self.post_minsize = post_minsize
-
         self.n_threads = n_threads
-
         self.dt_watershed = partial(dt_watershed,
                                     threshold=ws_threshold, sigma_seeds=ws_sigma,
                                     stacked=ws_2D, sigma_weights=ws_w_sigma,

plantseg/viewer/containers.py

@@ -13,6 +13,7 @@
 from plantseg.viewer.widget.proofreading.proofreading import widget_split_and_merge_from_scribbles
 from plantseg.viewer.widget.segmentation import widget_dt_ws, widget_agglomeration
 from plantseg.viewer.widget.segmentation import widget_fix_over_under_segmentation_from_nuclei
+from plantseg.viewer.widget.segmentation import widget_fix_false_positive_from_foreground_pmap
 from plantseg.viewer.widget.segmentation import widget_lifted_multicut
 from plantseg.viewer.widget.segmentation import widget_simple_dt_ws
 
@@ -65,7 +66,8 @@ def get_gasp_workflow():
 def get_extra_seg():
     container = MainWindow(widgets=[widget_dt_ws,
                                     widget_lifted_multicut,
-                                    widget_fix_over_under_segmentation_from_nuclei],
+                                    widget_fix_over_under_segmentation_from_nuclei,
+                                    widget_fix_false_positive_from_foreground_pmap],
                            labels=False)
     container = setup_menu(container, path='https://hci-unihd.github.io/plant-seg/chapters/plantseg_interactive_napari/extra_seg.html')
     return container

plantseg/viewer/widget/segmentation.py

@@ -7,7 +7,7 @@
 from napari.types import LayerDataTuple
 
 from napari import Viewer
-from plantseg.dataprocessing.functional.advanced_dataprocessing import fix_over_under_segmentation_from_nuclei
+from plantseg.dataprocessing.functional.advanced_dataprocessing import fix_over_under_segmentation_from_nuclei, remove_false_positives_by_foreground_probability
 from plantseg.dataprocessing.functional.dataprocessing import normalize_01
 from plantseg.segmentation.functional import gasp, multicut, dt_watershed, mutex_ws
 from plantseg.segmentation.functional import lifted_multicut_from_nuclei_segmentation, lifted_multicut_from_nuclei_pmaps
@@ -330,5 +330,36 @@ def widget_fix_over_under_segmentation_from_nuclei(cell_segmentation: Labels,
                                    input_keys=inputs_names,
                                    layer_kwarg=layer_kwargs,
                                    layer_type=layer_type,
-                                   step_name=f'Fix Over / Under segmentation',
+                                   step_name='Fix Over / Under Segmentation',
+                                   )
+
+
+@magicgui(call_button='Run Segmentation Fix from Foreground Pmap',
+          segmentation={'label': 'Segmentation'},
+          foreground={'label': 'Foreground Pmap'},
+          threshold={'label': 'Threshold',
+                     'widget_type': 'FloatSlider', 'max': 1., 'min': 0.})
+def widget_fix_false_positive_from_foreground_pmap(segmentation: Labels,
+                                                   foreground: Image,  # TODO: maybe also allow labels
+                                                   threshold=0.6) -> Future[LayerDataTuple]:
+    out_name = create_layer_name(segmentation.name, 'FGPmapFix')
+
+    inputs_names = (segmentation.name, foreground.name)
+    func_kwargs = {'segmentation': segmentation.data,
+                   'foreground': foreground.data}
+
+    layer_kwargs = layer_properties(name=out_name,
+                                    scale=segmentation.scale,
+                                    metadata=segmentation.metadata)
+    layer_type = 'labels'
+    step_kwargs = dict(threshold=threshold)
+
+    return start_threading_process(remove_false_positives_by_foreground_probability,
+                                   runtime_kwargs=func_kwargs,
+                                   statics_kwargs=step_kwargs,
+                                   out_name=out_name,
+                                   input_keys=inputs_names,
+                                   layer_kwarg=layer_kwargs,
+                                   layer_type=layer_type,
+                                   step_name='Reduce False Positives',
                                    )