Finalize CryoVesNet evaluation

Author: constantinpape

scripts/baselines/cryo_ves_net/README.md

@@ -1,5 +1,3 @@
 # CryoVesNet
 
-Scripts to run CryoVesNet on our data. See https://github.com/Zuber-group/CryoVesNet for details.
-
-The code is currently not working due to this issue: https://github.com/Zuber-group/CryoVesNet/issues/6
+Scripts to run CryoVesNet on our data and evalute the results. See https://github.com/Zuber-group/CryoVesNet for details.

scripts/baselines/cryo_ves_net/evaluate_cooper.py

@@ -1,12 +1,15 @@
+import json
 import os
 from glob import glob
 
 import h5py
 import numpy as np
 import pandas as pd
 from elf.evaluation.matching import matching
+from tqdm import tqdm
 
 INPUT_ROOT = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/vesicles_processed_v2/testsets"  # noqa
+INPUT_04 = "/mnt/lustre-emmy-hdd/projects/nim00007/data/synaptic-reconstruction/cooper/ground_truth/04Dataset_for_vesicle_eval"  # noqa
 OUTPUT_ROOT = "./predictions/cooper"  # noqa
 
 DATASETS = [
@@ -31,8 +34,17 @@ def evaluate_dataset(ds_name):
         return results
 
     print("Evaluating ds", ds_name)
-    input_files = sorted(glob(os.path.join(INPUT_ROOT, ds_name, "**/*.h5"), recursive=True))
+    if ds_name == "04":
+        input_files = sorted(glob(os.path.join(INPUT_04, "**/*.h5"), recursive=True))
+        seg_key = "labels/vesicles"
+        mask_key = "labels/compartment"
+    else:
+        input_files = sorted(glob(os.path.join(INPUT_ROOT, ds_name, "**/*.h5"), recursive=True))
+        seg_key = "/labels/vesicles/combined_vesicles"
+        mask_key = None
+
     pred_files = sorted(glob(os.path.join(OUTPUT_ROOT, ds_name, "**/*.h5"), recursive=True))
+    assert len(input_files) == len(pred_files), f"{len(input_files)}, {len(pred_files)}"
 
     results = {
         "dataset": [],
@@ -41,16 +53,45 @@ def evaluate_dataset(ds_name):
         "recall": [],
         "f1-score": [],
     }
-    for inf, predf in zip(input_files, pred_files):
+    for inf, predf in tqdm(zip(input_files, pred_files), total=len(input_files), desc=f"Evaluate {ds_name}"):
         fname = os.path.basename(inf)
-
-        with h5py.File(inf, "r") as f:
-            gt = f["/labels/vesicles/combined_vesicles"][:]
-        with h5py.File(predf, "r") as f:
-            seg = f["/prediction/vesicles/cryovesnet"][:]
-        assert gt.shape == seg.shape
-
-        scores = matching(seg, gt)
+        sub_res_path = os.path.join(result_folder, f"{ds_name}_{fname}.json")
+
+        if os.path.exists(sub_res_path):
+            print("Loading scores from", sub_res_path)
+            with open(sub_res_path, "r") as f:
+                scores = json.load(f)
+
+        else:
+            try:
+                with h5py.File(predf, "r") as f:
+                    seg = f["/prediction/vesicles/cryovesnet"][:]
+            except Exception:
+                print("Skipping", predf)
+                continue
+
+            with h5py.File(inf, "r") as f:
+                gt = f[seg_key][:]
+                if mask_key is None:
+                    mask = None
+                else:
+                    mask = f[mask_key][:]
+
+            assert gt.shape == seg.shape
+
+            if mask is not None:
+                bb = np.where(mask != 0)
+                bb = tuple(slice(
+                    int(b.min()), int(b.max()) + 1
+                ) for b in bb)
+                seg, gt, mask = seg[bb], gt[bb], mask[bb]
+                seg[mask == 0] = 0
+                gt[mask == 0] = 0
+
+            scores = matching(seg, gt)
+
+            with open(sub_res_path, "w") as f:
+                json.dump(scores, f)
 
         results["dataset"].append(ds_name)
         results["file"].append(fname)