torch_em.util.segmentation

  1from typing import Optional, List
  2
  3import numpy as np
  4
  5from skimage.feature import peak_local_max
  6
  7import bioimage_cpp as bic
  8
  9import elf.segmentation as elseg
 10from elf.segmentation.utils import normalize_input
 11from elf.segmentation.mutex_watershed import mutex_watershed
 12
 13
 14#
 15# Segmentation Functionality
 16#
 17
 18
 19def size_filter(
 20    seg: np.ndarray, min_size: int, hmap: Optional[np.ndarray] = None, with_background: bool = False
 21) -> np.ndarray:
 22    """Apply size filter to a segmentation to remove small segments.
 23
 24    Args:
 25        seg: The input segmentation.
 26        min_size: The minimal segmentation size.
 27        hmap: A heightmap to use for watershed based segmentation filtering.
 28        with_background: Whether this is a segmentation problem with background.
 29
 30    Returns:
 31        The filtered segmentation.
 32    """
 33    if min_size == 0:
 34        return seg
 35
 36    if hmap is None:
 37        ids, sizes = np.unique(seg, return_counts=True)
 38        bg_ids = ids[sizes < min_size]
 39        seg[np.isin(seg, bg_ids)] = 0
 40        seg, _, _ = bic.segmentation.relabel_sequential(seg.astype(np.uint), offset=1)
 41    else:
 42        assert hmap.ndim in (seg.ndim, seg.ndim + 1)
 43        hmap_ = np.max(hmap[:seg.ndim], axis=0) if hmap.ndim > seg.ndim else hmap
 44        if with_background:
 45            seg, _ = elseg.watershed.apply_size_filter(seg + 1, hmap_, min_size, exclude=[1])
 46            seg[seg == 1] = 0
 47        else:
 48            seg, _ = elseg.watershed.apply_size_filter(seg, hmap_, min_size)
 49    return seg
 50
 51
 52def mutex_watershed_segmentation(
 53    foreground: np.ndarray,
 54    affinities: np.ndarray,
 55    offsets: List[List[int]],
 56    min_size: int = 50,
 57    threshold: float = 0.5,
 58    strides: Optional[List[int]] = None
 59) -> np.ndarray:
 60    """Compute the mutex watershed segmentation using the affinity map for given pixel offsets.
 61
 62    Args:
 63        foreground: The foreground/background probabilities.
 64        affinities: The input affinity maps.
 65        offsets: The pixel offsets corresponding to the affinity channels.
 66        min_size: The minimum pixel size for objects in the output segmentation.
 67        threshold: The threshold for the foreground predictions.
 68        strides: The strides used to subsample long range edges.
 69
 70    Returns:
 71        The instance segmentation.
 72    """
 73    mask = (foreground >= threshold)
 74    if strides is None:
 75        strides = [2] * foreground.ndim
 76
 77    seg = mutex_watershed(affinities, offsets=offsets, mask=mask, strides=strides, randomize_strides=True)
 78    seg = size_filter(seg.astype("uint32"), min_size=min_size, hmap=affinities, with_background=True)
 79
 80    return seg
 81
 82
 83def connected_components_with_boundaries(
 84    foreground: np.ndarray, boundaries: np.ndarray, threshold: float = 0.5
 85) -> np.ndarray:
 86    """Compute instance segmentation based on foreground and boundary predictions.
 87
 88    Args:
 89        foreground: The foreground probability predictions.
 90        boundaries: The boundary probability predictions.
 91        threshold: The threshold for finding connected components.
 92
 93    Returns:
 94        The instance segmentation.
 95    """
 96    input_ = np.clip(foreground - boundaries, 0, 1)
 97    seeds = bic.segmentation.label(input_ > threshold)
 98    mask = normalize_input(foreground > threshold)
 99    seg = bic.segmentation.watershed(boundaries, markers=seeds, mask=mask.astype(bool))
100    return seg.astype("uint64")
101
102
103def watershed_from_components(
104    boundaries: np.ndarray,
105    foreground: np.ndarray,
106    min_size: int = 50,
107    threshold1: float = 0.5,
108    threshold2: float = 0.5,
109) -> np.ndarray:
110    """Compute an instance segmentation based on boundary and foreground predictions
111
112    The segmentation is computed as follows:
113    - Subtract the boundaries from the foreground to separate touching objects.
114    - Use the connected components of the result as seeds.
115    - Use the thresholded foreground predictions as mask to grow back the pieces
116      lost by subtracting the boundary prediction.
117
118    Args:
119        boundaries: The boundary probability predictions.
120        foreground: The foreground probability predictions.
121        min_size: The minimum pixel size for objects in the output segmentation.
122        threshold1: The threshold for finding connected components.
123        threshold2: The threshold for growing components via watershed on the boundary predictions.
124
125    Returns:
126        The instance segmentation.
127    """
128    seeds = bic.segmentation.label((foreground - boundaries) > threshold1)
129    mask = foreground > threshold2
130    seg = bic.segmentation.watershed(boundaries, seeds, mask=mask)
131    seg = size_filter(seg, min_size)
132    return seg
133
134
135def watershed_from_maxima(
136    boundaries: np.ndarray,
137    foreground: np.ndarray,
138    min_distance: int,
139    min_size: int = 50,
140    sigma: float = 1.0,
141    threshold1: float = 0.5,
142) -> np.ndarray:
143    """Compute an instance segmentation based on a seeded watershed from distance maxima.
144
145    This function thresholds the boundary probabilities, computes a distance transform,
146    finds the distance maxima, and then applies a seeded watershed to obtain the instance segmentation.
147    Compared to `watershed_from_components` this has the advantage that objects can be better separated,
148    but it may over-segment objects with complex shapes.
149
150    The min_distance parameter controls the minimal distance between seeds, which
151    corresponds to the minimal distance between object centers.
152
153    Args:
154        boundaries: The boundary probability predictions.
155        foreground: The foreground probability predictions.
156        min_size: The minimum pixel size for objects in the output segmentation.
157        min_distance: The minimum distance between peaks, see `from skimage.feature.peak_local_max`.
158        sigma: The standard deviation for smoothing the distance map before computing maxima.
159        threshold1: The threshold for foreground predictions.
160
161    Returns:
162        The instance segmentation.
163    """
164    mask = foreground > threshold1
165    boundary_distances = bic.distance.distance_transform(boundaries < 0.1)
166    boundary_distances[~mask] = 0
167    boundary_distances = bic.filters.gaussian_smoothing(boundary_distances, sigma)
168    seed_points = peak_local_max(boundary_distances, min_distance=min_distance, exclude_border=False)
169    seeds = np.zeros(mask.shape, dtype="uint32")
170    seeds[seed_points[:, 0], seed_points[:, 1]] = np.arange(1, len(seed_points) + 1)
171    seg = bic.segmentation.watershed(boundaries, markers=seeds, mask=foreground.astype(bool))
172    return size_filter(seg, min_size)
173
174
175def watershed_from_center_and_boundary_distances(
176    center_distances: np.ndarray,
177    boundary_distances: np.ndarray,
178    foreground_map: np.ndarray,
179    center_distance_threshold: float = 0.5,
180    boundary_distance_threshold: float = 0.5,
181    foreground_threshold: float = 0.5,
182    distance_smoothing: float = 1.6,
183    min_size: int = 0,
184    debug: bool = False,
185) -> np.ndarray:
186    """Compute an instance segmentation via a seeded watershed on distances to object centers and boundaries.
187
188    The seeds are computed by finding connected components where both distance predictions
189    are smaller than the respective thresholds. Using both distances is supposed to prevent merging
190    narrow adjacent objects (if only using the center distance) or finding multiple seeds for non-convex
191    cells (if only using the boundary distances).
192
193    Args:
194        center_distances: Distance prediction to the objcet centers.
195        boundary_distances: Inverted distance prediction to object boundaries.
196        foreground_map: Prediction for foreground probabilities.
197        center_distance_threshold: Center distance predictions below this value will be
198            used to find seeds (intersected with thresholded boundary distance predictions).
199        boundary_distance_threshold: oundary distance predictions below this value will be
200            used to find seeds (intersected with thresholded center distance predictions).
201        foreground_threshold: Foreground predictions above this value will be used as foreground mask.
202        distance_smoothing: Sigma value for smoothing the distance predictions.
203        min_size: Minimal object size in the segmentation result.
204        debug: Return all intermediate results in a dictionary for debugging.
205
206    Returns:
207        The instance segmentation.
208    """
209    if distance_smoothing > 0:
210        center_distances = bic.filters.gaussian_smoothing(center_distances, distance_smoothing)
211        boundary_distances = bic.filters.gaussian_smoothing(boundary_distances, distance_smoothing)
212
213    fg_mask = foreground_map > foreground_threshold
214
215    marker_map = np.logical_and(
216        center_distances < center_distance_threshold, boundary_distances < boundary_distance_threshold
217    )
218    marker_map[~fg_mask] = 0
219    markers = bic.segmentation.label(marker_map)
220
221    seg = bic.segmentation.watershed(boundary_distances, markers=markers, mask=fg_mask)
222    seg = size_filter(seg, min_size)
223
224    if debug:
225        debug_output = {
226            "center_distances": center_distances,
227            "boundary_distances": boundary_distances,
228            "foreground_mask": fg_mask,
229            "markers": markers,
230        }
231        return seg, debug_output
232
233    return seg