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