synapse_net.inference.vesicles

  1import time
  2import warnings
  3from typing import Dict, List, Optional, Tuple, Union
  4
  5import elf.parallel as parallel
  6import numpy as np
  7
  8import torch
  9
 10from synapse_net.inference.util import apply_size_filter, get_prediction, _Scaler
 11from synapse_net.inference.postprocessing.vesicles import filter_border_objects, filter_border_vesicles
 12from skimage.segmentation import relabel_sequential
 13
 14
 15def distance_based_vesicle_segmentation(
 16    foreground: np.ndarray,
 17    boundaries: np.ndarray,
 18    verbose: bool,
 19    min_size: int,
 20    boundary_threshold: float = 0.5,  # previous default value was 0.9
 21    distance_threshold: int = 8,
 22    block_shape: Tuple[int, int, int] = (128, 256, 256),
 23    halo: Tuple[int, int, int] = (48, 48, 48),
 24) -> np.ndarray:
 25    """Segment vesicles using a seeded watershed from connected components derived from
 26    distance transform of the boundary predictions.
 27
 28    This approach can prevent false merges that occur with the `simple_vesicle_segmentation`.
 29
 30    Args:
 31        foreground: The foreground prediction.
 32        boundaries: The boundary prediction.
 33        verbose: Whether to print timing information.
 34        min_size: The minimal vesicle size.
 35        boundary_threshold: The threshold for binarizing the boundary predictions for the distance computation.
 36        distance_threshold: The threshold for finding connected components in the boundary distances.
 37        block_shape: Block shape for parallelizing the operations.
 38        halo: Halo for parallelizing the operations.
 39
 40    Returns:
 41        The vesicle segmentation.
 42    """
 43    # Compute the boundary distances.
 44    t0 = time.time()
 45    bd_dist = parallel.distance_transform(
 46        boundaries < boundary_threshold, halo=halo, verbose=verbose, block_shape=block_shape
 47    )
 48    bd_dist[foreground < 0.5] = 0
 49    if verbose:
 50        print("Compute distance transform in", time.time() - t0, "s")
 51
 52    # Get the segmentation via seeded watershed of components in the boundary distances.
 53    t0 = time.time()
 54    seeds = parallel.label(bd_dist > distance_threshold, block_shape=block_shape, verbose=verbose)
 55    if verbose:
 56        print("Compute connected components in", time.time() - t0, "s")
 57
 58    # Compute distances from the seeds, which are used as heightmap for the watershed,
 59    # to assign all pixels to the nearest seed.
 60    t0 = time.time()
 61    dist = parallel.distance_transform(seeds == 0, halo=halo, verbose=verbose, block_shape=block_shape)
 62    if verbose:
 63        print("Compute distance transform in", time.time() - t0, "s")
 64
 65    t0 = time.time()
 66    mask = (foreground + boundaries) > 0.5
 67    seg = np.zeros_like(seeds)
 68    seg = parallel.seeded_watershed(
 69        dist, seeds, block_shape=block_shape,
 70        out=seg, mask=mask, verbose=verbose, halo=halo,
 71    )
 72    if verbose:
 73        print("Compute watershed in", time.time() - t0, "s")
 74
 75    seg = apply_size_filter(seg, min_size, verbose, block_shape)
 76    return seg
 77
 78
 79def simple_vesicle_segmentation(
 80    foreground: np.ndarray,
 81    boundaries: np.ndarray,
 82    verbose: bool,
 83    min_size: int,
 84    block_shape: Tuple[int, int, int] = (128, 256, 256),
 85    halo: Tuple[int, int, int] = (48, 48, 48),
 86) -> np.ndarray:
 87    """Segment vesicles by subtracting boundary from foreground prediction and
 88    applying connected components.
 89
 90    Args:
 91        foreground: The foreground prediction.
 92        boundaries: The boundary prediction.
 93        verbose: Whether to print timing information.
 94        min_size: The minimal vesicle size.
 95        block_shape: Block shape for parallelizing the operations.
 96        halo: Halo for parallelizing the operations.
 97
 98    Returns:
 99        The vesicle segmentation.
100    """
101
102    t0 = time.time()
103    seeds = parallel.label((foreground - boundaries) > 0.5, block_shape=block_shape, verbose=verbose)
104    if verbose:
105        print("Compute connected components in", time.time() - t0, "s")
106
107    t0 = time.time()
108    dist = parallel.distance_transform(seeds == 0, halo=halo, verbose=verbose, block_shape=block_shape)
109    if verbose:
110        print("Compute distance transform in", time.time() - t0, "s")
111
112    t0 = time.time()
113    mask = (foreground + boundaries) > 0.5
114    seg = np.zeros_like(seeds)
115    seg = parallel.seeded_watershed(
116        dist, seeds, block_shape=block_shape,
117        out=seg, mask=mask, verbose=verbose, halo=halo,
118    )
119    if verbose:
120        print("Compute watershed in", time.time() - t0, "s")
121
122    seg = apply_size_filter(seg, min_size, verbose, block_shape)
123    return seg
124
125
126def segment_vesicles(
127    input_volume: np.ndarray,
128    model_path: Optional[str] = None,
129    model: Optional[torch.nn.Module] = None,
130    tiling: Optional[Dict[str, Dict[str, int]]] = None,
131    min_size: int = 500,
132    verbose: bool = True,
133    distance_based_segmentation: bool = True,
134    return_predictions: bool = False,
135    scale: Optional[List[float]] = None,
136    exclude_boundary: bool = False,
137    exclude_boundary_vesicles: bool = False,
138    mask: Optional[np.ndarray] = None,
139) -> Union[np.ndarray, Tuple[np.ndarray, np.ndarray]]:
140    """Segment vesicles in an input volume or image.
141
142    Args:
143        input_volume: The input volume to segment.
144        model_path: The path to the model checkpoint if `model` is not provided.
145        model: Pre-loaded model. Either `model_path` or `model` is required.
146        tiling: The tiling configuration for the prediction.
147        min_size: The minimum size of a vesicle to be considered.
148        verbose: Whether to print timing information.
149        distance_based_segmentation: Whether to use distance-based segmentation.
150        return_predictions: Whether to return the predictions (foreground, boundaries) alongside the segmentation.
151        scale: The scale factor to use for rescaling the input volume before prediction.
152        exclude_boundary: Whether to exclude vesicles that touch the upper / lower border in z.
153        exclude_boundary_vesicles: Whether to exlude vesicles on the boundary that have less than the full diameter
154            inside of the volume. This is an alternative to post-processing with `exclude_boundary` that filters
155            out less vesicles at the boundary and is better suited for volumes with small context in z.
156            If `exclude_boundary` is also set to True, then this option will have no effect.
157        mask: An optional mask that is used to restrict the segmentation.
158
159    Returns:
160        The segmentation mask as a numpy array, or a tuple containing the segmentation mask
161        and the predictions if return_predictions is True.
162    """
163    if verbose:
164        print("Segmenting vesicles in volume of shape", input_volume.shape)
165    # Create the scaler to handle prediction with a different scaling factor.
166    scaler = _Scaler(scale, verbose)
167    input_volume = scaler.scale_input(input_volume)
168
169    # Rescale the mask if it was given and run prediction.
170    if mask is not None:
171        mask = scaler.scale_input(mask, is_segmentation=True)
172    pred = get_prediction(input_volume, tiling=tiling, model_path=model_path, model=model, verbose=verbose, mask=mask)
173    foreground, boundaries = pred[:2]
174
175    # Deal with 2D segmentation case.
176    kwargs = {}
177    if len(input_volume.shape) == 2:
178        kwargs["block_shape"] = (256, 256)
179        kwargs["halo"] = (48, 48)
180
181    if distance_based_segmentation:
182        seg = distance_based_vesicle_segmentation(
183            foreground, boundaries, verbose=verbose, min_size=min_size, **kwargs
184        )
185    else:
186        seg = simple_vesicle_segmentation(
187            foreground, boundaries, verbose=verbose, min_size=min_size, **kwargs
188        )
189
190    if exclude_boundary and exclude_boundary_vesicles:
191        warnings.warn(
192            "You have set both 'exclude_boundary' and 'exclude_boundary_vesicles' to True."
193            "The 'exclude_boundary_vesicles' option will have no effect."
194        )
195        seg = filter_border_objects(seg)
196
197    elif exclude_boundary:
198        seg = filter_border_objects(seg)
199
200    elif exclude_boundary_vesicles:
201        # Filter the vesicles that are at the z-border with less than their full diameter.
202        seg_ids = filter_border_vesicles(seg)
203
204        # Remove everything not in seg ids and relable the remaining IDs consecutively.
205        seg[~np.isin(seg, seg_ids)] = 0
206        seg = relabel_sequential(seg)[0]
207
208    seg = scaler.rescale_output(seg, is_segmentation=True)
209
210    if return_predictions:
211        pred = scaler.rescale_output(pred, is_segmentation=False)
212        return seg, pred
213    return seg