synapse_net.inference.vesicles

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