torch_em.util.segmentation
1import numpy as np 2 3import vigra 4import elf.segmentation as elseg 5from elf.segmentation.utils import normalize_input 6from elf.segmentation.mutex_watershed import mutex_watershed 7 8from skimage.measure import label 9from skimage.filters import gaussian 10from skimage.segmentation import watershed 11from skimage.feature import peak_local_max 12from scipy.ndimage import distance_transform_edt 13 14 15# 16# segmentation functionality 17# 18 19 20# could also refactor this into elf 21def size_filter(seg, min_size, hmap=None, with_background=False): 22 if min_size == 0: 23 return seg 24 25 if hmap is None: 26 ids, sizes = np.unique(seg, return_counts=True) 27 bg_ids = ids[sizes < min_size] 28 seg[np.isin(seg, bg_ids)] = 0 29 seg, _, _ = vigra.analysis.relabelConsecutive(seg.astype(np.uint), start_label=1, keep_zeros=True) 30 else: 31 assert hmap.ndim in (seg.ndim, seg.ndim + 1) 32 hmap_ = np.max(hmap[:seg.ndim], axis=0) if hmap.ndim > seg.ndim else hmap 33 if with_background: 34 seg, _ = elseg.watershed.apply_size_filter(seg + 1, hmap_, min_size, exclude=[1]) 35 seg[seg == 1] = 0 36 else: 37 seg, _ = elseg.watershed.apply_size_filter(seg, hmap_, min_size) 38 return seg 39 40 41def mutex_watershed_segmentation(foreground, affinities, offsets, min_size=50, threshold=0.5, strides=None): 42 """Computes the mutex watershed segmentation using the affinity maps for respective pixel offsets 43 44 Arguments: 45 - foreground: [np.ndarray] - The foreground background channel for the objects 46 - affinities [np.ndarray] - The input affinity maps 47 - offsets: [list[list[int]]] - The pixel offsets corresponding to the affinity channels 48 - min_size: [int] - The minimum pixels (below which) to filter objects 49 - threshold: [float] - To threshold foreground predictions 50 """ 51 mask = (foreground >= threshold) 52 if strides is None: 53 strides = [2] * foreground.ndim 54 seg = mutex_watershed(affinities, offsets=offsets, mask=mask, strides=strides, randomize_strides=True) 55 seg = size_filter(seg.astype("uint32"), min_size=min_size, hmap=affinities, with_background=True) 56 return seg 57 58 59def connected_components_with_boundaries(foreground, boundaries, threshold=0.5): 60 input_ = np.clip(foreground - boundaries, 0, 1) 61 seeds = label(input_ > threshold) 62 mask = normalize_input(foreground > threshold) 63 seg = watershed(boundaries, markers=seeds, mask=mask) 64 return seg.astype("uint64") 65 66 67def watershed_from_components(boundaries, foreground, min_size=50, threshold1=0.5, threshold2=0.5): 68 """The default approach: 69 - Subtract the boundaries from the foreground to separate touching objects. 70 - Use the connected components of this as seeds. 71 - Use the thresholded foreground predictions as mask to grow back the pieces 72 lost by subtracting the boundary prediction. 73 74 Arguments: 75 - boundaries: [np.ndarray] - The boundaries for objects 76 - foreground: [np.ndarray] - The foregrounds for objects 77 - min_size: [int] - The minimum pixels (below which) to filter objects 78 - threshold1: [float] - To separate touching objects (by subtracting bd and fg) above threshold 79 - threshold2: [float] - To threshold foreground predictions 80 81 Returns: 82 seg: [np.ndarray] - instance segmentation 83 """ 84 seeds = label((foreground - boundaries) > threshold1) 85 mask = foreground > threshold2 86 seg = watershed(boundaries, seeds, mask=mask) 87 seg = size_filter(seg, min_size) 88 return seg 89 90 91def watershed_from_maxima(boundaries, foreground, min_distance, min_size=50, sigma=1.0, threshold1=0.5): 92 """Find objects via seeded watershed starting from the maxima of the distance transform instead. 93 This has the advantage that objects can be better separated, but it may over-segment 94 if the objects have complex shapes. 95 96 The min_distance parameter controls the minimal distance between seeds, which 97 corresponds to the minimal distance between object centers. 98 99 Arguments: 100 - boundaries: [np.ndarray] - The boundaries for objects 101 - foreground: [np.ndarray] - The foreground for objects 102 - min_size: [int] - min. pixels (below which) to filter objects 103 - min_distance: [int] - min. distance of peaks (see `from skimage.feature import peak_local_max`) 104 - sigma: [float] - standard deviation for gaussian kernel. (see `from skimage.filters import gaussian`) 105 - threshold1: [float] - To threshold foreground predictions 106 107 Returns 108 seg: [np.ndarray] - instance segmentation 109 """ 110 mask = foreground > threshold1 111 boundary_distances = distance_transform_edt(boundaries < 0.1) 112 boundary_distances[~mask] = 0 # type: ignore 113 boundary_distances = gaussian(boundary_distances, sigma) # type: ignore 114 seed_points = peak_local_max(boundary_distances, min_distance=min_distance, exclude_border=False) 115 seeds = np.zeros(mask.shape, dtype="uint32") 116 seeds[seed_points[:, 0], seed_points[:, 1]] = np.arange(1, len(seed_points) + 1) 117 seg = watershed(boundaries, markers=seeds, mask=foreground) 118 seg = size_filter(seg, min_size) 119 return seg 120 121 122def watershed_from_center_and_boundary_distances( 123 center_distances, 124 boundary_distances, 125 foreground_map, 126 center_distance_threshold=0.5, 127 boundary_distance_threshold=0.5, 128 foreground_threshold=0.5, 129 distance_smoothing=1.6, 130 min_size=0, 131 debug=False, 132): 133 """Seeded watershed based on distance predictions to object center and boundaries. 134 135 The seeds are computed by finding connected components where both distance predictions 136 are smaller than the respective thresholds. Using both distances here should prevent merging 137 narrow adjacent objects (if only using the center distance) or finding multiple seeds for non-convex 138 cells (if only using the boundary distances). 139 140 Args: 141 center_distances [np.ndarray] - Distance prediction to the objcet center. 142 boundary_distances [np.ndarray] - Inverted distance prediction to object boundaries. 143 foreground_map [np.ndarray] - Predictio for foreground probabilities. 144 center_distance_threshold [float] - Center distance predictions below this value will be 145 used to find seeds (intersected with thresholded boundary distance predictions). 146 boundary_distance_threshold [float] - Boundary distance predictions below this value will be 147 used to find seeds (intersected with thresholded center distance predictions). 148 foreground_threshold [float] - Foreground predictions above this value will be used as foreground mask. 149 distance_smoothing [float] - Sigma value for smoothing the distance predictions. 150 min_size [int] - Minimal object size in the segmentation result. 151 debug [bool] - Return all intermediate results for debugging. 152 153 Returns: 154 np.ndarray - The instance segmentation. 155 """ 156 center_distances = vigra.filters.gaussianSmoothing(center_distances, distance_smoothing) 157 boundary_distances = vigra.filters.gaussianSmoothing(boundary_distances, distance_smoothing) 158 159 fg_mask = foreground_map > foreground_threshold 160 161 marker_map = np.logical_and( 162 center_distances < center_distance_threshold, 163 boundary_distances < boundary_distance_threshold 164 ) 165 marker_map[~fg_mask] = 0 166 markers = label(marker_map) 167 168 seg = watershed(boundary_distances, markers=markers, mask=fg_mask) 169 seg = size_filter(seg, min_size) 170 171 if debug: 172 debug_output = { 173 "center_distances": center_distances, 174 "boundary_distances": boundary_distances, 175 "foreground_mask": fg_mask, 176 "markers": markers, 177 } 178 return seg, debug_output 179 180 return seg
22def size_filter(seg, min_size, hmap=None, with_background=False): 23 if min_size == 0: 24 return seg 25 26 if hmap is None: 27 ids, sizes = np.unique(seg, return_counts=True) 28 bg_ids = ids[sizes < min_size] 29 seg[np.isin(seg, bg_ids)] = 0 30 seg, _, _ = vigra.analysis.relabelConsecutive(seg.astype(np.uint), start_label=1, keep_zeros=True) 31 else: 32 assert hmap.ndim in (seg.ndim, seg.ndim + 1) 33 hmap_ = np.max(hmap[:seg.ndim], axis=0) if hmap.ndim > seg.ndim else hmap 34 if with_background: 35 seg, _ = elseg.watershed.apply_size_filter(seg + 1, hmap_, min_size, exclude=[1]) 36 seg[seg == 1] = 0 37 else: 38 seg, _ = elseg.watershed.apply_size_filter(seg, hmap_, min_size) 39 return seg
42def mutex_watershed_segmentation(foreground, affinities, offsets, min_size=50, threshold=0.5, strides=None): 43 """Computes the mutex watershed segmentation using the affinity maps for respective pixel offsets 44 45 Arguments: 46 - foreground: [np.ndarray] - The foreground background channel for the objects 47 - affinities [np.ndarray] - The input affinity maps 48 - offsets: [list[list[int]]] - The pixel offsets corresponding to the affinity channels 49 - min_size: [int] - The minimum pixels (below which) to filter objects 50 - threshold: [float] - To threshold foreground predictions 51 """ 52 mask = (foreground >= threshold) 53 if strides is None: 54 strides = [2] * foreground.ndim 55 seg = mutex_watershed(affinities, offsets=offsets, mask=mask, strides=strides, randomize_strides=True) 56 seg = size_filter(seg.astype("uint32"), min_size=min_size, hmap=affinities, with_background=True) 57 return seg
Computes the mutex watershed segmentation using the affinity maps for respective pixel offsets
Arguments:
- - foreground: [np.ndarray] - The foreground background channel for the objects
- affinities [np.ndarray] - The input affinity maps
- - offsets: [list[list[int]]] - The pixel offsets corresponding to the affinity channels
- - min_size: [int] - The minimum pixels (below which) to filter objects
- - threshold: [float] - To threshold foreground predictions
60def connected_components_with_boundaries(foreground, boundaries, threshold=0.5): 61 input_ = np.clip(foreground - boundaries, 0, 1) 62 seeds = label(input_ > threshold) 63 mask = normalize_input(foreground > threshold) 64 seg = watershed(boundaries, markers=seeds, mask=mask) 65 return seg.astype("uint64")
68def watershed_from_components(boundaries, foreground, min_size=50, threshold1=0.5, threshold2=0.5): 69 """The default approach: 70 - Subtract the boundaries from the foreground to separate touching objects. 71 - Use the connected components of this as seeds. 72 - Use the thresholded foreground predictions as mask to grow back the pieces 73 lost by subtracting the boundary prediction. 74 75 Arguments: 76 - boundaries: [np.ndarray] - The boundaries for objects 77 - foreground: [np.ndarray] - The foregrounds for objects 78 - min_size: [int] - The minimum pixels (below which) to filter objects 79 - threshold1: [float] - To separate touching objects (by subtracting bd and fg) above threshold 80 - threshold2: [float] - To threshold foreground predictions 81 82 Returns: 83 seg: [np.ndarray] - instance segmentation 84 """ 85 seeds = label((foreground - boundaries) > threshold1) 86 mask = foreground > threshold2 87 seg = watershed(boundaries, seeds, mask=mask) 88 seg = size_filter(seg, min_size) 89 return seg
The default approach:
- Subtract the boundaries from the foreground to separate touching objects.
- Use the connected components of this as seeds.
- Use the thresholded foreground predictions as mask to grow back the pieces lost by subtracting the boundary prediction.
Arguments:
- - boundaries: [np.ndarray] - The boundaries for objects
- - foreground: [np.ndarray] - The foregrounds for objects
- - min_size: [int] - The minimum pixels (below which) to filter objects
- - threshold1: [float] - To separate touching objects (by subtracting bd and fg) above threshold
- - threshold2: [float] - To threshold foreground predictions
Returns:
seg: [np.ndarray] - instance segmentation
92def watershed_from_maxima(boundaries, foreground, min_distance, min_size=50, sigma=1.0, threshold1=0.5): 93 """Find objects via seeded watershed starting from the maxima of the distance transform instead. 94 This has the advantage that objects can be better separated, but it may over-segment 95 if the objects have complex shapes. 96 97 The min_distance parameter controls the minimal distance between seeds, which 98 corresponds to the minimal distance between object centers. 99 100 Arguments: 101 - boundaries: [np.ndarray] - The boundaries for objects 102 - foreground: [np.ndarray] - The foreground for objects 103 - min_size: [int] - min. pixels (below which) to filter objects 104 - min_distance: [int] - min. distance of peaks (see `from skimage.feature import peak_local_max`) 105 - sigma: [float] - standard deviation for gaussian kernel. (see `from skimage.filters import gaussian`) 106 - threshold1: [float] - To threshold foreground predictions 107 108 Returns 109 seg: [np.ndarray] - instance segmentation 110 """ 111 mask = foreground > threshold1 112 boundary_distances = distance_transform_edt(boundaries < 0.1) 113 boundary_distances[~mask] = 0 # type: ignore 114 boundary_distances = gaussian(boundary_distances, sigma) # type: ignore 115 seed_points = peak_local_max(boundary_distances, min_distance=min_distance, exclude_border=False) 116 seeds = np.zeros(mask.shape, dtype="uint32") 117 seeds[seed_points[:, 0], seed_points[:, 1]] = np.arange(1, len(seed_points) + 1) 118 seg = watershed(boundaries, markers=seeds, mask=foreground) 119 seg = size_filter(seg, min_size) 120 return seg
Find objects via seeded watershed starting from the maxima of the distance transform instead. This has the advantage that objects can be better separated, but it may over-segment if the objects have complex shapes.
The min_distance parameter controls the minimal distance between seeds, which corresponds to the minimal distance between object centers.
Arguments:
- - boundaries: [np.ndarray] - The boundaries for objects
- - foreground: [np.ndarray] - The foreground for objects
- - min_size: [int] - min. pixels (below which) to filter objects
- - min_distance: [int] - min. distance of peaks (see
from skimage.feature import peak_local_max
) - - sigma: [float] - standard deviation for gaussian kernel. (see
from skimage.filters import gaussian
) - - threshold1: [float] - To threshold foreground predictions
Returns seg: [np.ndarray] - instance segmentation
123def watershed_from_center_and_boundary_distances( 124 center_distances, 125 boundary_distances, 126 foreground_map, 127 center_distance_threshold=0.5, 128 boundary_distance_threshold=0.5, 129 foreground_threshold=0.5, 130 distance_smoothing=1.6, 131 min_size=0, 132 debug=False, 133): 134 """Seeded watershed based on distance predictions to object center and boundaries. 135 136 The seeds are computed by finding connected components where both distance predictions 137 are smaller than the respective thresholds. Using both distances here should prevent merging 138 narrow adjacent objects (if only using the center distance) or finding multiple seeds for non-convex 139 cells (if only using the boundary distances). 140 141 Args: 142 center_distances [np.ndarray] - Distance prediction to the objcet center. 143 boundary_distances [np.ndarray] - Inverted distance prediction to object boundaries. 144 foreground_map [np.ndarray] - Predictio for foreground probabilities. 145 center_distance_threshold [float] - Center distance predictions below this value will be 146 used to find seeds (intersected with thresholded boundary distance predictions). 147 boundary_distance_threshold [float] - Boundary distance predictions below this value will be 148 used to find seeds (intersected with thresholded center distance predictions). 149 foreground_threshold [float] - Foreground predictions above this value will be used as foreground mask. 150 distance_smoothing [float] - Sigma value for smoothing the distance predictions. 151 min_size [int] - Minimal object size in the segmentation result. 152 debug [bool] - Return all intermediate results for debugging. 153 154 Returns: 155 np.ndarray - The instance segmentation. 156 """ 157 center_distances = vigra.filters.gaussianSmoothing(center_distances, distance_smoothing) 158 boundary_distances = vigra.filters.gaussianSmoothing(boundary_distances, distance_smoothing) 159 160 fg_mask = foreground_map > foreground_threshold 161 162 marker_map = np.logical_and( 163 center_distances < center_distance_threshold, 164 boundary_distances < boundary_distance_threshold 165 ) 166 marker_map[~fg_mask] = 0 167 markers = label(marker_map) 168 169 seg = watershed(boundary_distances, markers=markers, mask=fg_mask) 170 seg = size_filter(seg, min_size) 171 172 if debug: 173 debug_output = { 174 "center_distances": center_distances, 175 "boundary_distances": boundary_distances, 176 "foreground_mask": fg_mask, 177 "markers": markers, 178 } 179 return seg, debug_output 180 181 return seg
Seeded watershed based on distance predictions to object center and boundaries.
The seeds are computed by finding connected components where both distance predictions are smaller than the respective thresholds. Using both distances here should 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 [np.ndarray] - Distance prediction to the objcet center.
- boundary_distances [np.ndarray] - Inverted distance prediction to object boundaries.
- foreground_map [np.ndarray] - Predictio for foreground probabilities.
- center_distance_threshold [float] - Center distance predictions below this value will be used to find seeds (intersected with thresholded boundary distance predictions).
- boundary_distance_threshold [float] - Boundary distance predictions below this value will be used to find seeds (intersected with thresholded center distance predictions).
- foreground_threshold [float] - Foreground predictions above this value will be used as foreground mask.
- distance_smoothing [float] - Sigma value for smoothing the distance predictions.
- min_size [int] - Minimal object size in the segmentation result.
- debug [bool] - Return all intermediate results for debugging.
Returns:
np.ndarray - The instance segmentation.