Implement scalable segmentatio WIP

Author: constantinpape

synapse_net/inference/scalable_segmentation.py

@@ -0,0 +1,99 @@
+import os
+import tempfile
+from typing import Dict, Optional
+
+import elf.parallel as parallel
+import numpy as np
+import torch
+
+from elf.io import open_file
+from elf.wrapper import ThresholdWrapper, SimpleTransformationWrapper
+from elf.wrapper.base import MultiTransformationWrapper
+from synapse_net.inference.util import get_prediction
+from numpy.typing import ArrayLike
+
+
+class SelectChannel(SimpleTransformationWrapper):
+    """Wrapper to select a chanel from an array-like dataset object.
+
+    Args:
+        volume: The array-like input dataset.
+        channel: The channel that will be selected.
+    """
+    def __init__(self, volume: np.typing.ArrayLike, channel: int):
+        self.channel = channel
+        super().__init__(volume, lambda x: x[self.channel], with_channels=True)
+
+    @property
+    def shape(self):
+        return self._volume.shape[1:]
+
+    @property
+    def chunks(self):
+        return self._volume.chunks[1:]
+
+    @property
+    def ndim(self):
+        return self._volume.ndim - 1
+
+
+# TODO support resizing via the wrapper
+def scalable_segmentation(
+    input_: ArrayLike,
+    output: ArrayLike,
+    model: torch.nn.Module,
+    tiling: Optional[Dict[str, Dict[str, int]]] = None,
+    seed_threshold: float = 0.5,
+    min_size: int = 500,
+    verbose: bool = True,
+) -> None:
+    """Lorem ipsum
+
+    Args:
+        input_:
+        output:
+        model: The model.
+        tiling: The tiling configuration for the prediction.
+        min_size: The minimum size of a vesicle to be considered.
+        verbose: Whether to print timing information.
+    """
+    assert model.out_channels == 2
+
+    # Create a temporary directory for storing the predictions.
+    with tempfile.TemporaryDirectory() as tmp_dir:
+        tmp_pred = os.path.join(tmp_dir, "prediction.n5")
+        f = open_file(tmp_pred, mode="a")
+
+        # Create the dataset for storing the prediction.
+        chunks = (128,) * 3
+        pred_shape = (2,) + input_.shape
+        pred_chunks = (1,) + chunks
+        pred = f.create_dataset("pred", shape=pred_shape, dtype="float32", chunks=pred_chunks)
+
+        # Run the prediction.
+        get_prediction(input_, prediction=pred, tiling=tiling, model=model, verbose=verbose)
+
+        # Create wrappers for selecting the foreground and the boundary channel.
+        foreground = SelectChannel(pred, 0)
+        boundaries = SelectChannel(pred, 1)
+
+        # Create temporary storage for the seeds.
+        tmp_seeds = os.path.join(tmp_dir, "seeds.n5")
+        f = open_file(tmp_seeds, mode="a")
+        seeds = f.create_dataset("seeds", shape=input_.shape, dtype="uint64", chunks=chunks)
+
+        # Create wrappers for subtracting and thresholding boundary subtracted from the foreground.
+        # And then compute the seeds based on this.
+        seed_input = ThresholdWrapper(
+            MultiTransformationWrapper(np.subtract, foreground, boundaries), seed_threshold
+        )
+        parallel.label(seed_input, seeds, verbose=verbose)
+
+        # Run watershed to extend back from the seeds to the boundaries.
+        parallel.seeded_watershed(
+            boundaries, seeds=seeds, out=output, verbose=verbose, block_shape=chunks, halo=3 * (16,)
+        )
+
+        # Run the size filter.
+        if min_size > 0:
+            parallel.size_filter(output, output, min_size=min_size, verbose=verbose)

synapse_net/inference/util.py

@@ -18,6 +18,7 @@
 # import xarray
 
 from elf.io import open_file
+from numpy.typing import ArrayLike
 from scipy.ndimage import binary_closing
 from skimage.measure import regionprops
 from skimage.morphology import remove_small_holes
@@ -100,15 +101,16 @@ def rescale_output(self, output, is_segmentation):
 
 
 def get_prediction(
-    input_volume: np.ndarray,  # [z, y, x]
+    input_volume: ArrayLike,  # [z, y, x]
     tiling: Optional[Dict[str, Dict[str, int]]],  # {"tile": {"z": int, ...}, "halo": {"z": int, ...}}
     model_path: Optional[str] = None,
     model: Optional[torch.nn.Module] = None,
     verbose: bool = True,
     with_channels: bool = False,
     channels_to_standardize: Optional[List[int]] = None,
-    mask: Optional[np.ndarray] = None,
-) -> np.ndarray:
+    mask: Optional[ArrayLike] = None,
+    prediction: Optional[ArrayLike] = None,
+) -> ArrayLike:
     """Run prediction on a given volume.
 
     This function will automatically choose the correct prediction implementation,
@@ -124,6 +126,8 @@ def get_prediction(
         channels_to_standardize: List of channels to standardize. Defaults to None.
         mask: Optional binary mask. If given, the prediction will only be run in
             the foreground region of the mask.
+        prediction: An array like object for writing the prediction.
+            If not given, the prediction will be computed in moemory.
 
     Returns:
         The predicted volume.
@@ -174,21 +178,23 @@ def get_prediction(
         for dim in tiling["tile"]:
             updated_tiling["tile"][dim] = tiling["tile"][dim] - 2 * tiling["halo"][dim]
         # print(f"updated_tiling {updated_tiling}")
-        pred = get_prediction_torch_em(
-            input_volume, updated_tiling, model_path, model, verbose, with_channels, mask=mask
+        prediction = get_prediction_torch_em(
+            input_volume, updated_tiling, model_path, model, verbose, with_channels,
+            mask=mask, prediction=prediction,
         )
 
-    return pred
+    return prediction
 
 
 def get_prediction_torch_em(
-    input_volume: np.ndarray,  # [z, y, x]
+    input_volume: ArrayLike,  # [z, y, x]
     tiling: Dict[str, Dict[str, int]],  # {"tile": {"z": int, ...}, "halo": {"z": int, ...}}
     model_path: Optional[str] = None,
     model: Optional[torch.nn.Module] = None,
     verbose: bool = True,
     with_channels: bool = False,
-    mask: Optional[np.ndarray] = None,
+    mask: Optional[ArrayLike] = None,
+    prediction: Optional[ArrayLike] = None,
 ) -> np.ndarray:
     """Run prediction using torch-em on a given volume.
 
@@ -201,6 +207,8 @@ def get_prediction_torch_em(
         with_channels: Whether to predict with channels.
         mask: Optional binary mask. If given, the prediction will only be run in
             the foreground region of the mask.
+        prediction: An array like object for writing the prediction.
+            If not given, the prediction will be computed in moemory.
 
     Returns:
         The predicted volume.
@@ -234,14 +242,15 @@ def get_prediction_torch_em(
                 print("Run prediction with mask.")
             mask = mask.astype("bool")
 
-        pred = predict_with_halo(
+        prediction = predict_with_halo(
             input_volume, model, gpu_ids=[device],
             block_shape=block_shape, halo=halo,
             preprocess=None, with_channels=with_channels, mask=mask,
+            output=prediction,
         )
     if verbose:
         print("Prediction time in", time.time() - t0, "s")
-    return pred
+    return prediction
 
 
 def _get_file_paths(input_path, ext=".mrc"):