Updates to support large-scale prediction with the synapse detection model

Author: constantinpape

flamingo_tools/segmentation/unet_prediction.py

@@ -8,6 +8,7 @@
 import numpy as np
 import nifty.tools as nt
 import vigra
+import tifffile
 import torch
 import z5py
 
@@ -37,7 +38,10 @@ def ndim(self):
         return self._volume.ndim - 1
 
 
-def prediction_impl(input_path, input_key, output_folder, model_path, scale, block_shape, halo):
+def prediction_impl(
+    input_path, input_key, output_folder, model_path, scale, block_shape, halo,
+    output_channels=3, apply_postprocessing=True,
+):
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
         if os.path.isdir(model_path):
@@ -46,10 +50,16 @@ def prediction_impl(input_path, input_key, output_folder, model_path, scale, blo
             model = torch.load(model_path)
 
     mask_path = os.path.join(output_folder, "mask.zarr")
-    image_mask = z5py.File(mask_path, "r")["mask"]
+    if os.path.exists(mask_path):
+        image_mask = z5py.File(mask_path, "r")["mask"]
+    else:
+        image_mask = None
 
     if input_key is None:
-        input_ = imageio.imread(input_path)
+        try:
+            input_ = tifffile.memmap(input_path)
+        except Exception:
+            input_ = imageio.imread(input_path)
     else:
         input_ = open_file(input_path, "r")[input_key]
 
@@ -93,17 +103,27 @@ def preprocess(raw):
         raw /= std
         return raw
 
-    # Smooth the distance prediction channel.
-    def postprocess(x):
-        x[1] = vigra.filters.gaussianSmoothing(x[1], sigma=2.0)
-        return x
+    if apply_postprocessing:
+        # Smooth the distance prediction channel.
+        def postprocess(x):
+            x[1] = vigra.filters.gaussianSmoothing(x[1], sigma=2.0)
+            return x
+    else:
+        postprocess = None if output_channels > 1 else lambda x: x.squeeze()
+
+    if output_channels > 1:
+        output_shape = (output_channels,) + input_.shape
+        output_chunks = (1,) + block_shape
+    else:
+        output_shape = input_.shape
+        output_chunks = block_shape
 
     output_path = os.path.join(output_folder, "predictions.zarr")
     with open_file(output_path, "a") as f:
         output = f.require_dataset(
             "prediction",
-            shape=(3,) + input_.shape,
-            chunks=(1,) + block_shape,
+            shape=output_shape,
+            chunks=output_chunks,
             compression="gzip",
             dtype="float32",
         )

scripts/synapse_marker_detection/run_prediction.py

@@ -0,0 +1,56 @@
+import argparse
+import os
+import sys
+
+import pandas as pd
+import numpy as np
+import zarr
+
+from elf.parallel.local_maxima import find_local_maxima
+
+sys.path.append("../..")
+
+
+def main():
+    from flamingo_tools.segmentation.unet_prediction import prediction_impl
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--input", required=True)
+    parser.add_argument("-o", "--output_folder", required=True)
+    parser.add_argument("-m", "--model", required=True)
+    parser.add_argument("-k", "--input_key", default=None)
+    args = parser.parse_args()
+
+    block_shape = (64, 256, 256)
+    halo = (16, 64, 64)
+
+    # Skip existing prediction, which is saved in output_folder/predictions.zarr
+    skip_prediction = False
+    output_path = os.path.join(args.output_folder, "predictions.zarr")
+    prediction_key = "prediction"
+    if os.path.exists(output_path) and prediction_key in zarr.open(output_path, "r"):
+        skip_prediction = True
+
+    if not skip_prediction:
+        prediction_impl(
+            args.input, args.input_key, args.output_folder, args.model,
+            scale=None, block_shape=block_shape, halo=halo,
+            apply_postprocessing=False, output_channels=1,
+        )
+
+    detection_path = os.path.join(args.output_folder, "synapse_detection.tsv")
+    if not os.path.exists(detection_path):
+        input_ = zarr.open(output_path, "r")[prediction_key]
+        detections = find_local_maxima(
+            input_, block_shape=block_shape, min_distance=2, threshold_abs=0.5, verbose=True, n_threads=16,
+        )
+        # Save the result in mobie compatible format.
+        detections = np.concatenate(
+            [np.arange(1, len(detections) + 1)[:, None], detections[:, ::-1]], axis=1
+        )
+        detections = pd.DataFrame(detections, columns=["spot_id", "x", "y", "z"])
+        detections.to_csv(detection_path, index=False, sep="\t")
+
+
+if __name__ == "__main__":
+    main()