torch_em.util.segmentation

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