Update segmentation CLI

Author: constantinpape

flamingo_tools/segmentation/cli.py

@@ -20,7 +20,7 @@ def run_segmentation():
     """
     parser = argparse.ArgumentParser(description="")
     parser.add_argument("-i", "--input_path", required=True, help="The path to the input data.")
-    parser.add_argument("-k", "--input_key", required=True, help="The key to the input data.")
+    parser.add_argument("-k", "--input_key", help="The key to the input data.")
     parser.add_argument("-o", "--output_folder", required=True)
     parser.add_argument("-m", "--model_type", required=True)
     parser.add_argument("-c", "--checkpoint_path")
@@ -43,14 +43,20 @@ def run_detection():
     """private
     """
     parser = argparse.ArgumentParser()
+    parser.add_argument("-i", "--input_path", required=True, help="The path to the input data.")
+    parser.add_argument("-k", "--input_key", help="The key to the input data.")
+    parser.add_argument("-o", "--output_folder", required=True)
     parser.add_argument("-m", "--model_type", default="Synapses")
+    parser.add_argument("--mask_path")
+    parser.add_argument("--mask_key")
+    parser.add_argument("-c", "--checkpoint_path")
     args = parser.parse_args()
     detection_models = ["Synapses"]
     if args.model_type not in detection_models:
         raise ValueError
     model_path = _get_model_path(args.model_type, args.checkpoint_path)
-    # TODO
     marker_detection(
         input_path=args.input_path, input_key=args.input_key,
         output_folder=args.output_folder, model_path=model_path,
+        mask_path=args.mask_path, mask_input_key=args.mask_key,
     )

flamingo_tools/segmentation/synapse_detection.py

@@ -119,10 +119,10 @@ def run_prediction(
 def marker_detection(
     input_path: str,
     input_key: str,
-    mask_path: str,
+    mask_path: Optional[str],
     output_folder: str,
     model_path: str,
-    mask_input_key: str = "s4",
+    mask_input_key: Optional[str] = "s4",
     max_distance: float = 20,
     resolution: float = 0.38,
 ):
@@ -143,13 +143,12 @@ def marker_detection(
     # Best approach: load IHC segmentation at a low scale level, binarize it,
     # dilate it and use this as mask. It can be mapped back to the full resolution
     # with `elf.wrapper.ResizedVolume`.
-
     skip_masking = False
 
     mask_preprocess_key = "mask"
     output_file = os.path.join(output_folder, "mask.zarr")
 
-    if os.path.exists(output_file) and mask_preprocess_key in zarr.open(output_file, "r"):
+    if mask_path is None or (os.path.exists(output_file) and mask_preprocess_key in zarr.open(output_file, "r")):
         skip_masking = True
 
     if not skip_masking:
@@ -162,11 +161,6 @@ def marker_detection(
             f_out.create_dataset(mask_preprocess_key, data=arr_bin, compression="gzip")
 
     # 2.) Run inference and detection of maxima.
-    # This can be taken from 'scripts/synapse_marker_detection/run_prediction.py'
-    # (And the run prediction script should then be refactored).
-
-    block_shape = (64, 256, 256)
-    halo = (16, 64, 64)
 
     # Skip existing prediction, which is saved in output_folder/predictions.zarr
     skip_prediction = False
@@ -183,11 +177,12 @@ def marker_detection(
     if not skip_prediction:
         prediction_impl(
             input_path, input_key, output_folder, model_path,
-            scale=None, block_shape=block_shape, halo=halo,
-            apply_postprocessing=False, output_channels=1,
+            scale=None, apply_postprocessing=False, output_channels=1,
+            block_shape=None, halo=None,
         )
 
     if not os.path.exists(detection_path):
+        block_shape = (64, 256, 256)
         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,
@@ -200,25 +195,25 @@ def marker_detection(
         detections.to_csv(detection_path, index=False, sep="\t")
 
     else:
-        with open(detection_path, 'r') as f:
+        with open(detection_path, "r") as f:
             detections = pd.read_csv(f, sep="\t")
 
     # 3.) Map the detections to IHC and filter them based on a distance criterion.
     # Use the function 'map_and_filter_detections' from above.
-    input_ = read_image_data(mask_path, input_key)
-
-    detections_filtered = map_and_filter_detections(
-        segmentation=input_,
-        detections=detections,
-        max_distance=max_distance,
-        resolution=resolution,
-    )
+    if mask_path is not None:
+        input_ = read_image_data(mask_path, input_key)
+        detections_filtered = map_and_filter_detections(
+            segmentation=input_,
+            detections=detections,
+            max_distance=max_distance,
+            resolution=resolution,
+        )
 
-    # 4.) Add the filtered detections to MoBIE.
-    # IMPORTANT scale the coordinates with the resolution here.
-    detections_filtered["distance_to_ihc"] *= resolution
-    detections_filtered["x"] *= resolution
-    detections_filtered["y"] *= resolution
-    detections_filtered["z"] *= resolution
-    detection_path = os.path.join(output_folder, "synapse_detection_filtered.tsv")
-    detections_filtered.to_csv(detection_path, index=False, sep="\t")
+        # 4.) Add the filtered detections to MoBIE.
+        # IMPORTANT scale the coordinates with the resolution here.
+        detections_filtered["distance_to_ihc"] *= resolution
+        detections_filtered["x"] *= resolution
+        detections_filtered["y"] *= resolution
+        detections_filtered["z"] *= resolution
+        detection_path = os.path.join(output_folder, "synapse_detection_filtered.tsv")
+        detections_filtered.to_csv(detection_path, index=False, sep="\t")

flamingo_tools/segmentation/unet_prediction.py

@@ -55,6 +55,21 @@ def ndim(self):
         return self._volume.ndim - 1
 
 
+def _get_device_and_tiling(block_shape, halo, input_):
+    have_cuda = torch.cuda.is_available()
+    if block_shape is None:
+        block_shape = (128, 128, 128) if have_cuda else getattr(input_, "chunks", (64, 64, 64))
+    if halo is None:
+        halo = (16, 32, 32)
+    if have_cuda:
+        print("Predict with GPU")
+        gpu_ids = [0]
+    else:
+        print("Predict with CPU")
+        gpu_ids = ["cpu"]
+    return gpu_ids, block_shape, halo
+
+
 def prediction_impl(
     input_path,
     input_key,
@@ -109,19 +124,6 @@ def prediction_impl(
         input_ = ResizedVolume(input_, shape=new_shape, order=3)
         image_mask = ResizedVolume(image_mask, new_shape, order=0)
 
-    have_cuda = torch.cuda.is_available()
-
-    if block_shape is None:
-        block_shape = (128, 128, 128) if have_cuda else input_.chunks
-    if halo is None:
-        halo = (16, 32, 32)
-    if have_cuda:
-        print("Predict with GPU")
-        gpu_ids = [0]
-    else:
-        print("Predict with CPU")
-        gpu_ids = ["cpu"]
-
     if mean is None or std is None:
         # Compute the global mean and standard deviation.
         n_threads = min(16, mp.cpu_count())
@@ -156,6 +158,7 @@ def postprocess(x):
 
     shape = input_.shape
     ndim = len(shape)
+    gpu_ids, block_shape, halo = _get_device_and_tiling(block_shape, halo, input_)
 
     blocking = nt.blocking([0] * ndim, shape, block_shape)
     n_blocks = blocking.numberOfBlocks

test/test_cli.py

@@ -19,9 +19,9 @@ def setUp(self):
     def tearDown(self):
         rmtree(self.folder)
 
-    def test_convert_flamingo(self):
+    def test_convert_data(self):
         out_path = os.path.join(self.folder, "converted_data.n5")
-        cmd = ["flamingo_tools.convert_flamingo", "-i", self.folder, "-o", out_path, "--metadata_pattern", ""]
+        cmd = ["flamingo_tools.convert_data", "-i", self.folder, "-o", out_path, "--metadata_pattern", ""]
         run(cmd)
 
         self.assertTrue(os.path.exists(out_path))

test/test_segmentation/test_cli.py

@@ -0,0 +1,52 @@
+import os
+import subprocess
+import tempfile
+import unittest
+
+import imageio.v3 as imageio
+import numpy as np
+import z5py
+
+
+class TestSegmentationCLI(unittest.TestCase):
+    shape = (64, 128, 128)
+
+    def _create_data(self, tmp_dir):
+        data = np.random.randint(0, 255, size=self.shape)
+        path = os.path.join(tmp_dir, "data.tif")
+        imageio.imwrite(path, data)
+        return path
+
+    def test_run_segmentation(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            data_path = self._create_data(tmp_dir)
+            output_folder = os.path.join(tmp_dir, "output")
+
+            subprocess.run([
+                "flamingo_tools.run_segmentation",
+                "-i", data_path, "-o", output_folder, "-m", "SGN", "--min_size", "0"
+            ])
+
+            expected_path = os.path.join(output_folder, "segmentation.zarr")
+            expected_key = "segmentation"
+
+            self.assertTrue(os.path.exists(expected_path))
+            with z5py.File(expected_path, "r") as f:
+                self.assertTrue(expected_key in f)
+                self.assertEqual(f[expected_key].shape, self.shape)
+
+    def test_run_detection(self):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            data_path = self._create_data(tmp_dir)
+            output_folder = os.path.join(tmp_dir, "output")
+
+            subprocess.run([
+                "flamingo_tools.run_detection", "-i", data_path, "-o", output_folder
+            ])
+
+            expected_path = os.path.join(output_folder, "synapse_detection.tsv")
+            self.assertTrue(os.path.exists(expected_path))
+
+
+if __name__ == "__main__":
+    unittest.main()