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
def size_filter( seg: numpy.ndarray, min_size: int, hmap: Optional[numpy.ndarray] = None, with_background: bool = False) -> numpy.ndarray:
24def size_filter(
25    seg: np.ndarray, min_size: int, hmap: Optional[np.ndarray] = None, with_background: bool = False
26) -> np.ndarray:
27    """Apply size filter to a segmentation to remove small segments.
28
29    Args:
30        seg: The input segmentation.
31        min_size: The minimal segmentation size.
32        hmap: A heightmap to use for watershed based segmentation filtering.
33        with_background: Whether this is a segmentation problem with background.
34
35    Returns:
36        The filtered segmentation.
37    """
38    if min_size == 0:
39        return seg
40
41    if hmap is None:
42        ids, sizes = np.unique(seg, return_counts=True)
43        bg_ids = ids[sizes < min_size]
44        seg[np.isin(seg, bg_ids)] = 0
45        seg, _, _ = vigra.analysis.relabelConsecutive(seg.astype(np.uint), start_label=1, keep_zeros=True)
46    else:
47        assert hmap.ndim in (seg.ndim, seg.ndim + 1)
48        hmap_ = np.max(hmap[:seg.ndim], axis=0) if hmap.ndim > seg.ndim else hmap
49        if with_background:
50            seg, _ = elseg.watershed.apply_size_filter(seg + 1, hmap_, min_size, exclude=[1])
51            seg[seg == 1] = 0
52        else:
53            seg, _ = elseg.watershed.apply_size_filter(seg, hmap_, min_size)
54    return seg

Apply size filter to a segmentation to remove small segments.

Arguments:
  • seg: The input segmentation.
  • min_size: The minimal segmentation size.
  • hmap: A heightmap to use for watershed based segmentation filtering.
  • with_background: Whether this is a segmentation problem with background.
Returns:

The filtered segmentation.

def mutex_watershed_segmentation( foreground: numpy.ndarray, affinities: numpy.ndarray, offsets: List[List[int]], min_size: int = 50, threshold: float = 0.5, strides: Optional[List[int]] = None) -> numpy.ndarray:
57def mutex_watershed_segmentation(
58    foreground: np.ndarray,
59    affinities: np.ndarray,
60    offsets: List[List[int]],
61    min_size: int = 50,
62    threshold: float = 0.5,
63    strides: Optional[List[int]] = None
64) -> np.ndarray:
65    """Compute the mutex watershed segmentation using the affinity map for given pixel offsets.
66
67    Args:
68        foreground: The foreground/background probabilities.
69        affinities: The input affinity maps.
70        offsets: The pixel offsets corresponding to the affinity channels.
71        min_size: The minimum pixel size for objects in the output segmentation.
72        threshold: The threshold for the foreground predictions.
73        strides: The strides used to subsample long range edges.
74
75    Returns:
76        The instance segmentation.
77    """
78    mask = (foreground >= threshold)
79    if strides is None:
80        strides = [2] * foreground.ndim
81
82    seg = mutex_watershed(affinities, offsets=offsets, mask=mask, strides=strides, randomize_strides=True)
83    seg = size_filter(seg.astype("uint32"), min_size=min_size, hmap=affinities, with_background=True)
84
85    return seg

Compute the mutex watershed segmentation using the affinity map for given pixel offsets.

Arguments:
  • foreground: The foreground/background probabilities.
  • affinities: The input affinity maps.
  • offsets: The pixel offsets corresponding to the affinity channels.
  • min_size: The minimum pixel size for objects in the output segmentation.
  • threshold: The threshold for the foreground predictions.
  • strides: The strides used to subsample long range edges.
Returns:

The instance segmentation.

def connected_components_with_boundaries( foreground: numpy.ndarray, boundaries: numpy.ndarray, threshold: float = 0.5) -> numpy.ndarray:
 88def connected_components_with_boundaries(
 89    foreground: np.ndarray, boundaries: np.ndarray, threshold: float = 0.5
 90) -> np.ndarray:
 91    """Compute instance segmentation based on foreground and boundary predictions.
 92
 93    Args:
 94        foreground: The foreground probability predictions.
 95        boundaries: The boundary probability predictions.
 96        threshold: The threshold for finding connected components.
 97
 98    Returns:
 99        The instance segmentation.
100    """
101    input_ = np.clip(foreground - boundaries, 0, 1)
102    seeds = label(input_ > threshold)
103    mask = normalize_input(foreground > threshold)
104    seg = watershed(boundaries, markers=seeds, mask=mask)
105    return seg.astype("uint64")

Compute instance segmentation based on foreground and boundary predictions.

Arguments:
  • foreground: The foreground probability predictions.
  • boundaries: The boundary probability predictions.
  • threshold: The threshold for finding connected components.
Returns:

The instance segmentation.

def watershed_from_components( boundaries: numpy.ndarray, foreground: numpy.ndarray, min_size: int = 50, threshold1: float = 0.5, threshold2: float = 0.5) -> numpy.ndarray:
108def watershed_from_components(
109    boundaries: np.ndarray,
110    foreground: np.ndarray,
111    min_size: int = 50,
112    threshold1: float = 0.5,
113    threshold2: float = 0.5,
114) -> np.ndarray:
115    """Compute an instance segmentation based on boundary and foreground predictions
116
117    The segmentation is computed as follows:
118    - Subtract the boundaries from the foreground to separate touching objects.
119    - Use the connected components of the result as seeds.
120    - Use the thresholded foreground predictions as mask to grow back the pieces
121      lost by subtracting the boundary prediction.
122
123    Args:
124        boundaries: The boundary probability predictions.
125        foreground: The foreground probability predictions.
126        min_size: The minimum pixel size for objects in the output segmentation.
127        threshold1: The threshold for finding connected components.
128        threshold2: The threshold for growing components via watershed on the boundary predictions.
129
130    Returns:
131        The instance segmentation.
132    """
133    seeds = label((foreground - boundaries) > threshold1)
134    mask = foreground > threshold2
135    seg = watershed(boundaries, seeds, mask=mask)
136    seg = size_filter(seg, min_size)
137    return seg

Compute an instance segmentation based on boundary and foreground predictions

The segmentation is computed as follows:

  • Subtract the boundaries from the foreground to separate touching objects.
  • Use the connected components of the result as seeds.
  • Use the thresholded foreground predictions as mask to grow back the pieces lost by subtracting the boundary prediction.
Arguments:
  • boundaries: The boundary probability predictions.
  • foreground: The foreground probability predictions.
  • min_size: The minimum pixel size for objects in the output segmentation.
  • threshold1: The threshold for finding connected components.
  • threshold2: The threshold for growing components via watershed on the boundary predictions.
Returns:

The instance segmentation.

def watershed_from_maxima( boundaries: numpy.ndarray, foreground: numpy.ndarray, min_distance: int, min_size: int = 50, sigma: float = 1.0, threshold1: float = 0.5) -> numpy.ndarray:
140def watershed_from_maxima(
141    boundaries: np.ndarray,
142    foreground: np.ndarray,
143    min_distance: int,
144    min_size: int = 50,
145    sigma: float = 1.0,
146    threshold1: float = 0.5,
147) -> np.ndarray:
148    """Compute an instance segmentation based on a seeded watershed from distance maxima.
149
150    This function thresholds the boundary probabilities, computes a distance transform,
151    finds the distance maxima, and then applies a seeded watershed to obtain the instance segmentation.
152    Compared to `watershed_from_components` this has the advantage that objects can be better separated,
153    but it may over-segment objects with complex shapes.
154
155    The min_distance parameter controls the minimal distance between seeds, which
156    corresponds to the minimal distance between object centers.
157
158    Args:
159        boundaries: The boundary probability predictions.
160        foreground: The foreground probability predictions.
161        min_size: The minimum pixel size for objects in the output segmentation.
162        min_distance: The minimum distance between peaks, see `from skimage.feature.peak_local_max`.
163        sigma: The standard deviation for smoothing the distance map before computing maxima.
164        threshold1: The threshold for foreground predictions.
165
166    Returns:
167        The instance segmentation.
168    """
169    mask = foreground > threshold1
170    boundary_distances = distance_transform_edt(boundaries < 0.1)
171    boundary_distances[~mask] = 0
172    boundary_distances = gaussian(boundary_distances, sigma)
173    seed_points = peak_local_max(boundary_distances, min_distance=min_distance, exclude_border=False)
174    seeds = np.zeros(mask.shape, dtype="uint32")
175    seeds[seed_points[:, 0], seed_points[:, 1]] = np.arange(1, len(seed_points) + 1)
176    seg = watershed(boundaries, markers=seeds, mask=foreground)
177    return size_filter(seg, min_size)

Compute an instance segmentation based on a seeded watershed from distance maxima.

This function thresholds the boundary probabilities, computes a distance transform, finds the distance maxima, and then applies a seeded watershed to obtain the instance segmentation. Compared to watershed_from_components this has the advantage that objects can be better separated, but it may over-segment objects with complex shapes.

The min_distance parameter controls the minimal distance between seeds, which corresponds to the minimal distance between object centers.

Arguments:
  • boundaries: The boundary probability predictions.
  • foreground: The foreground probability predictions.
  • min_size: The minimum pixel size for objects in the output segmentation.
  • min_distance: The minimum distance between peaks, see from skimage.feature.peak_local_max.
  • sigma: The standard deviation for smoothing the distance map before computing maxima.
  • threshold1: The threshold for foreground predictions.
Returns:

The instance segmentation.

def watershed_from_center_and_boundary_distances( center_distances: numpy.ndarray, boundary_distances: numpy.ndarray, foreground_map: numpy.ndarray, center_distance_threshold: float = 0.5, boundary_distance_threshold: float = 0.5, foreground_threshold: float = 0.5, distance_smoothing: float = 1.6, min_size: int = 0, debug: bool = False) -> numpy.ndarray:
180def watershed_from_center_and_boundary_distances(
181    center_distances: np.ndarray,
182    boundary_distances: np.ndarray,
183    foreground_map: np.ndarray,
184    center_distance_threshold: float = 0.5,
185    boundary_distance_threshold: float = 0.5,
186    foreground_threshold: float = 0.5,
187    distance_smoothing: float = 1.6,
188    min_size: int = 0,
189    debug: bool = False,
190) -> np.ndarray:
191    """Compute an instance segmentation via a seeded watershed on distances to object centers and boundaries.
192
193    The seeds are computed by finding connected components where both distance predictions
194    are smaller than the respective thresholds. Using both distances is supposed to prevent merging
195    narrow adjacent objects (if only using the center distance) or finding multiple seeds for non-convex
196    cells (if only using the boundary distances).
197
198    Args:
199        center_distances: Distance prediction to the objcet centers.
200        boundary_distances: Inverted distance prediction to object boundaries.
201        foreground_map: Prediction for foreground probabilities.
202        center_distance_threshold: Center distance predictions below this value will be
203            used to find seeds (intersected with thresholded boundary distance predictions).
204        boundary_distance_threshold: oundary distance predictions below this value will be
205            used to find seeds (intersected with thresholded center distance predictions).
206        foreground_threshold: Foreground predictions above this value will be used as foreground mask.
207        distance_smoothing: Sigma value for smoothing the distance predictions.
208        min_size: Minimal object size in the segmentation result.
209        debug: Return all intermediate results in a dictionary for debugging.
210
211    Returns:
212        The instance segmentation.
213    """
214    center_distances = vigra.filters.gaussianSmoothing(center_distances, distance_smoothing)
215    boundary_distances = vigra.filters.gaussianSmoothing(boundary_distances, distance_smoothing)
216
217    fg_mask = foreground_map > foreground_threshold
218
219    marker_map = np.logical_and(
220        center_distances < center_distance_threshold, boundary_distances < boundary_distance_threshold
221    )
222    marker_map[~fg_mask] = 0
223    markers = label(marker_map)
224
225    seg = watershed(boundary_distances, markers=markers, mask=fg_mask)
226    seg = size_filter(seg, min_size)
227
228    if debug:
229        debug_output = {
230            "center_distances": center_distances,
231            "boundary_distances": boundary_distances,
232            "foreground_mask": fg_mask,
233            "markers": markers,
234        }
235        return seg, debug_output
236
237    return seg

Compute an instance segmentation via a seeded watershed on distances to object centers and boundaries.

The seeds are computed by finding connected components where both distance predictions are smaller than the respective thresholds. Using both distances is supposed to prevent merging narrow adjacent objects (if only using the center distance) or finding multiple seeds for non-convex cells (if only using the boundary distances).

Arguments:
  • center_distances: Distance prediction to the objcet centers.
  • boundary_distances: Inverted distance prediction to object boundaries.
  • foreground_map: Prediction for foreground probabilities.
  • center_distance_threshold: Center distance predictions below this value will be used to find seeds (intersected with thresholded boundary distance predictions).
  • boundary_distance_threshold: oundary distance predictions below this value will be used to find seeds (intersected with thresholded center distance predictions).
  • foreground_threshold: Foreground predictions above this value will be used as foreground mask.
  • distance_smoothing: Sigma value for smoothing the distance predictions.
  • min_size: Minimal object size in the segmentation result.
  • debug: Return all intermediate results in a dictionary for debugging.
Returns:

The instance segmentation.