Update synapse prediction script

Author: constantinpape

scripts/synapse_marker_detection/check_synapse_prediction.py

@@ -1,10 +1,17 @@
+import os
+from glob import glob
+
 import h5py
+import imageio.v3 as imageio
 import napari
 import zarr
 
 from torch_em.util import load_model
 from torch_em.util.prediction import predict_with_halo
 from train_synapse_detection import get_paths
+from tqdm import tqdm
+
+OUTPUT_ROOT = "./predictions"
 
 
 def run_prediction(val_image):
@@ -15,22 +22,56 @@ def run_prediction(val_image):
     return pred.squeeze()
 
 
-def main():
-    val_paths, _ = get_paths("val")
-    val_image = zarr.open(val_paths[0])["raw"][:]
-
-    # pred = run_prediction(val_image)
-    # with h5py.File("pred.h5", "a") as f:
-    #     f.create_dataset("pred", data=pred, compression="gzip")
+def require_prediction(image_data, output_path):
+    key = "prediction"
+    if os.path.exists(output_path):
+        with h5py.File(output_path, "r") as f:
+            pred = f[key][:]
+    else:
+        pred = run_prediction(image_data)
+        with h5py.File(output_path, "w") as f:
+            f.create_dataset(key, data=pred, compression="gzip")
+    return pred
 
-    with h5py.File("pred.h5", "r") as f:
-        pred = f["pred"][:]
 
+def visualize_results(image_data, pred):
     v = napari.Viewer()
-    v.add_image(val_image)
+    v.add_image(image_data)
     v.add_image(pred)
     napari.run()
 
 
+def check_val_image():
+    val_paths, _ = get_paths("val")
+    val_path = val_paths[0]
+    val_image = zarr.open(val_path)["raw"][:]
+
+    os.makedirs(os.path.join(OUTPUT_ROOT, "val"), exist_ok=True)
+    output_path = os.path.join(OUTPUT_ROOT, "val", os.path.basename(val_path).replace(".zarr", ".h5"))
+    pred = require_prediction(val_image, output_path)
+
+    visualize_results(val_image, pred)
+
+
+def check_new_images():
+    input_root = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/synapses/test_crops"
+    inputs = glob(os.path.join(input_root, "*.tif"))
+    output_folder = os.path.join(OUTPUT_ROOT, "new_crops")
+    os.makedirs(output_folder, exist_ok=True)
+    for path in tqdm(inputs):
+        name = os.path.basename(path)
+        if name == "M_AMD_58L_avgblendfused_RibB.tif":
+            continue
+        image_data = imageio.imread(path)
+        output_path = os.path.join(output_folder, name.replace(".tif", ".h5"))
+        require_prediction(image_data, output_path)
+
+
+# TODO update to support post-processing and showing annotations for the val data
+def main():
+    # check_val_image()
+    check_new_images()
+
+
 if __name__ == "__main__":
     main()

scripts/synapse_marker_detection/train_synapse_detection.py

@@ -3,13 +3,14 @@
 
 from detection_dataset import DetectionDataset
 
-sys.path.append("/home/pape/Work/my_projects/czii-protein-challenge")
-# sys.path.append("/user/pape41/u12086/Work/my_projects/czii-protein-challenge")
+# sys.path.append("/home/pape/Work/my_projects/czii-protein-challenge")
+sys.path.append("/user/pape41/u12086/Work/my_projects/czii-protein-challenge")
 
 from utils.training.training import supervised_training  # noqa
 
-TRAIN_ROOT = "./training_data/images"
-LABEL_ROOT = "./training_data/labels"
+ROOT = "/mnt/vast-nhr/projects/nim00007/data/moser/cochlea-lightsheet/training_data/synapses/training_data/v1"  # noqa
+TRAIN_ROOT = os.path.join(ROOT, "images")
+LABEL_ROOT = os.path.join(ROOT, "labels")
 
 
 def get_paths(split):