Update model check script

Author: constantinpape

development/check_uploaded_models.py

@@ -8,6 +8,7 @@
 import zarr
 from flamingo_tools.test_data import _sample_registry
 
+view = False
 data_dict = {
     "SGN": "PV",
     "IHC": "VGlut3",
@@ -25,17 +26,19 @@ def check_segmentation_model(model_name):
         data_path = _sample_registry().fetch(data_dict[model_name])
         copyfile(data_path, input_path)
 
-    subprocess.run(
-        ["flamingo_tools.run_segmentation", "-i", input_path, "-o", output_folder, "-m", model_name]
-    )
     output_path = os.path.join(output_folder, "segmentation.zarr")
-    segmentation = zarr.open(output_path)["segmentation"][:]
+    if not os.path.exists(output_path):
+        subprocess.run(
+            ["flamingo_tools.run_segmentation", "-i", input_path, "-o", output_folder, "-m", model_name]
+        )
 
-    image = imageio.imread(input_path)
-    v = napari.Viewer()
-    v.add_image(image)
-    v.add_labels(segmentation, name=f"{model_name}-segmentation")
-    napari.run()
+    if view:
+        segmentation = zarr.open(output_path)["segmentation"][:]
+        image = imageio.imread(input_path)
+        v = napari.Viewer()
+        v.add_image(image)
+        v.add_labels(segmentation, name=f"{model_name}-segmentation")
+        napari.run()
 
 
 def check_detection_model():
@@ -47,36 +50,38 @@ def check_detection_model():
         data_path = _sample_registry().fetch(data_dict[model_name])
         copyfile(data_path, input_path)
 
-    subprocess.run(
-        ["flamingo_tools.run_detection", "-i", input_path, "-o", output_folder, "-m", model_name]
-    )
     output_path = os.path.join(output_folder, "synapse_detection.tsv")
-    prediction = pd.read_csv(output_path, sep="\t")[["z", "y", "x"]]
+    if not os.path.exists(output_path):
+        subprocess.run(
+            ["flamingo_tools.run_detection", "-i", input_path, "-o", output_folder, "-m", model_name]
+        )
 
-    image = imageio.imread(input_path)
-    v = napari.Viewer()
-    v.add_image(image)
-    v.add_points(prediction)
-    napari.run()
+    if view:
+        prediction = pd.read_csv(output_path, sep="\t")[["z", "y", "x"]]
+        image = imageio.imread(input_path)
+        v = napari.Viewer()
+        v.add_image(image)
+        v.add_points(prediction)
+        napari.run()
 
 
 def main():
     # SGN segmentation:
     # - Prediction works well on the CPU.
-    # check_segmentation_model("SGN")
+    check_segmentation_model("SGN")
 
     # IHC segmentation:
     # - Prediction does not work well on the CPU.
-    # check_segmentation_model("IHC")
+    check_segmentation_model("IHC")
 
     # SGN segmentation (lowres):
     # - Prediction does not work well on the CPU.
-    # check_segmentation_model("SGN-lowres")
+    check_segmentation_model("SGN-lowres")
 
     # IHC segmentation (lowres):
     # - The prediction seems to work (on the CPU), but a lot of merges.
     # -> Update the segmentation params?
-    # check_segmentation_model("IHC-lowres")
+    check_segmentation_model("IHC-lowres")
 
     # Synapse detection:
     # - Prediction works well on the CPU.