Add segmentation and measurement scripts for SGN stainings

Author: constantinpape

flamingo_tools/segmentation/unet_prediction.py

@@ -324,6 +324,7 @@ def run_unet_prediction(
     scale: Optional[float] = None,
     block_shape: Optional[Tuple[int, int, int]] = None,
     halo: Optional[Tuple[int, int, int]] = None,
+    use_mask: bool = True,
 ) -> None:
     """Run prediction and segmentation with a distance U-Net.
 
@@ -337,10 +338,12 @@ def run_unet_prediction(
             By default the data will not be rescaled.
         block_shape: The block-shape for running the prediction.
         halo: The halo (= block overlap) to use for prediction.
+        use_mask: Whether to use the masking heuristics to not run inference on empty blocks.
     """
     os.makedirs(output_folder, exist_ok=True)
 
-    find_mask(input_path, input_key, output_folder)
+    if use_mask:
+        find_mask(input_path, input_key, output_folder)
 
     original_shape = prediction_impl(
         input_path, input_key, output_folder, model_path, scale, block_shape, halo

scripts/sgn_stain_predictions/check_segmentation.py

@@ -0,0 +1,35 @@
+import os
+from glob import glob
+
+import imageio.v3 as imageio
+import napari
+
+
+ROOT = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/LS_sampleprepcomparison_crops"
+SAVE_ROOT = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/LS_sampleprepcomparison_crops/segmentations"
+
+
+def main():
+    files = sorted(glob(os.path.join(ROOT, "**/*.tif")))
+    for ff in files:
+        if "segmentations" in ff:
+            return
+        print("Visualizing", ff)
+        rel_path = os.path.relpath(ff, ROOT)
+        seg_path = os.path.join(SAVE_ROOT, rel_path)
+
+        image = imageio.imread(ff)
+        if os.path.exists(seg_path):
+            seg = imageio.imread(seg_path)
+        else:
+            seg = None
+
+        v = napari.Viewer()
+        v.add_image(image)
+        if seg is not None:
+            v.add_labels(seg)
+        napari.run()
+
+
+if __name__ == "__main__":
+    main()

scripts/sgn_stain_predictions/measure_intensities.py

@@ -0,0 +1,39 @@
+import os
+from glob import glob
+
+import tifffile
+from flamingo_tools.measurements import compute_object_measures_impl
+
+
+ROOT = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/LS_sampleprepcomparison_crops"
+SAVE_ROOT = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/LS_sampleprepcomparison_crops/segmentations"
+
+
+def measure_intensities(ff):
+    rel_path = os.path.relpath(ff, ROOT)
+    out_path = os.path.join("./measurements", rel_path.replace(".tif", ".xlsx"))
+    if os.path.exists(out_path):
+        return
+
+    print("Computing measurements for", rel_path)
+    seg_path = os.path.join(SAVE_ROOT, rel_path)
+
+    image_data = tifffile.memmap(ff)
+    seg_data = tifffile.memmap(seg_path)
+
+    table = compute_object_measures_impl(image_data, seg_data, n_threads=8)
+
+    os.makedirs(os.path.split(out_path)[0], exist_ok=True)
+    table.to_excel(out_path, index=False)
+
+
+def main():
+    files = sorted(glob(os.path.join(ROOT, "**/*.tif")))
+    for ff in files:
+        if "segmentations" in ff:
+            return
+        measure_intensities(ff)
+
+
+if __name__ == "__main__":
+    main()

scripts/sgn_stain_predictions/run_prediction.py

@@ -0,0 +1,58 @@
+import os
+import tempfile
+from glob import glob
+
+import tifffile
+from elf.io import open_file
+from flamingo_tools.segmentation import run_unet_prediction
+
+ROOT = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/LS_sampleprepcomparison_crops"
+MODEL_PATH = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/trained_models/SGN/cochlea_distance_unet_SGN_March2025Model"  # noqa
+
+SAVE_ROOT = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/LS_sampleprepcomparison_crops/segmentations"
+
+
+def check_data():
+    files = glob(os.path.join(ROOT, "**/*.tif"), recursive=True)
+    for ff in files:
+        rel_path = sorted(os.path.relpath(ff, ROOT))
+        shape = tifffile.memmap(ff).shape
+        print(rel_path, shape)
+
+
+def segment_crop(input_file):
+    fname = os.path.relpath(input_file, ROOT)
+    out_file = os.path.join(SAVE_ROOT, fname)
+    if "segmentations" in input_file:
+        return
+    if os.path.exists(out_file):
+        return
+
+    print("Run prediction for", input_file)
+    os.makedirs(os.path.split(out_file)[0], exist_ok=True)
+    with tempfile.TemporaryDirectory() as tmp_folder:
+        run_unet_prediction(
+            input_file, input_key=None, output_folder=tmp_folder,
+            model_path=MODEL_PATH, min_size=1000, use_mask=False,
+        )
+        seg_path = os.path.join(tmp_folder, "segmentation.zarr")
+        with open_file(seg_path, mode="r") as f:
+            seg = f["segmentation"][:]
+
+    print("Writing output to", out_file)
+    tifffile.imwrite(out_file, seg, bigtiff=True)
+
+
+def segment_all():
+    files = sorted(glob(os.path.join(ROOT, "**/*.tif"), recursive=True))
+    for ff in files:
+        segment_crop(ff)
+
+
+def main():
+    # check_data()
+    segment_all()
+
+
+if __name__ == "__main__":
+    main()