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
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.
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.
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.
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.
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.
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.