Extend generalist evaluation and fix some issues in evaluation logic

Author: constantinpape

finetuning/generalists/compile_results.py

@@ -0,0 +1,71 @@
+import os
+from glob import glob
+
+import pandas as pd
+
+from evaluate_generalist import EXPERIMENT_ROOT
+from util import EM_DATASETS, LM_DATASETS
+
+
+def get_results(model, ds):
+    res_folder = os.path.join(EXPERIMENT_ROOT, model, ds, "results")
+    res_paths = sorted(glob(os.path.join(res_folder, "box", "*.csv"))) +\
+        sorted(glob(os.path.join(res_folder, "points", "*.csv")))
+
+    amg_res = os.path.join(res_folder, "amg.csv")
+    if os.path.exists(amg_res):
+        res_paths.append(amg_res)
+
+    results = []
+    for path in res_paths:
+        prompt_res = pd.read_csv(path)
+        prompt_name = os.path.splitext(os.path.relpath(path, res_folder))[0]
+        prompt_res.insert(0, "prompt", [prompt_name])
+        results.append(prompt_res)
+    results = pd.concat(results)
+    results.insert(0, "dataset", results.shape[0] * [ds])
+
+    return results
+
+
+def compile_results(models, datasets, out_path):
+    results = []
+
+    for model in models:
+        model_results = []
+
+        for ds in datasets:
+            ds_results = get_results(model, ds)
+            model_results.append(ds_results)
+
+        model_results = pd.concat(model_results)
+        model_results.insert(0, "model", [model] * model_results.shape[0])
+        results.append(model_results)
+
+    results = pd.concat(results)
+    results.to_csv(out_path, index=False)
+
+
+def compile_em():
+    compile_results(
+        ["vit_h", "vit_h_em", "vit_b", "vit_b_em"],
+        EM_DATASETS,
+        os.path.join(EXPERIMENT_ROOT, "evaluation-em.csv")
+    )
+
+
+# TODO
+def compile_lm():
+    compile_results(
+        ["vit_h", "vit_h_lm", "vit_b", "vit_b_lm"],
+        LM_DATASETS,
+        os.path.join(EXPERIMENT_ROOT, "evaluation-lm.csv")
+    )
+
+
+def main():
+    compile_em()
+
+
+if __name__ == "__main__":
+    main()

finetuning/generalists/evaluate_generalist.py

@@ -11,11 +11,11 @@
     "vit_b": "/home/nimcpape/.sam_models/sam_vit_b_01ec64.pth",
     "vit_h": "/home/nimcpape/.sam_models/sam_vit_h_4b8939.pth",
     # Generalist LM models
-    "vit_b_lm": "/scratch-grete/projects/nim00007/sam/models/LM/generalist/v2/vit_b/best.pt",
-    "vit_h_lm": "/scratch-grete/projects/nim00007/sam/models/LM/generalist/v2/vit_h/best.pt",
+    "vit_b_lm": "/scratch/projects/nim00007/sam/models/LM/generalist/v2/vit_b/best.pt",
+    "vit_h_lm": "/scratch/projects/nim00007/sam/models/LM/generalist/v2/vit_h/best.pt",
     # Generalist EM models
-    "vit_b_em": "/scratch-grete/projects/nim00007/sam/models/EM/generalist/v2/vit_b/best.pt",
-    "vit_h_em": "/scratch-grete/projects/nim00007/sam/models/EM/generalist/v2/vit_h/best.pt",
+    "vit_b_em": "/scratch/projects/nim00007/sam/models/EM/generalist/v2/vit_b/best.pt",
+    "vit_h_em": "/scratch/projects/nim00007/sam/models/EM/generalist/v2/vit_h/best.pt",
 }
 
 

finetuning/generalists/export_generalist_model.py

@@ -0,0 +1,24 @@
+import os
+from micro_sam.util import export_custom_sam_model
+from evaluate_generalist import CHECKPOINTS, EXPERIMENT_ROOT
+
+OUT_ROOT = os.path.join(EXPERIMENT_ROOT, "exported")
+os.makedirs(OUT_ROOT, exist_ok=True)
+
+
+def export_generalist(model):
+    checkpoint_path = CHECKPOINTS[model]
+    model_type = model[:5]
+    save_path = os.path.join(OUT_ROOT, f"{model}.pth")
+    export_custom_sam_model(checkpoint_path, model_type, save_path)
+
+
+def main():
+    export_generalist("vit_b_em")
+    export_generalist("vit_h_em")
+    export_generalist("vit_b_lm")
+    export_generalist("vit_h_lm")
+
+
+if __name__ == "__main__":
+    main()

micro_sam/evaluation/evaluation.py

@@ -7,6 +7,7 @@
 import pandas as pd
 
 from elf.evaluation import mean_segmentation_accuracy
+from skimage.measure import label
 from tqdm import tqdm
 
 
@@ -21,6 +22,7 @@ def _run_evaluation(gt_paths, prediction_paths, verbose=True):
         assert os.path.exists(pred_path), pred_path
 
         gt = imageio.imread(gt_path)
+        gt = label(gt)
         pred = imageio.imread(pred_path)
 
         msa, scores = mean_segmentation_accuracy(pred, gt, return_accuracies=True)

micro_sam/util.py

@@ -501,8 +501,8 @@ def precompute_image_embeddings(
                 # key signature does not match or is not in the file
                 if key not in f.attrs or f.attrs[key] != val:
                     warnings.warn(
-                        f"Embeddings file {save_path} is invalid due to unmatching {key}. "
-                        "Please recompute embeddings in a new file."
+                        f"Embeddings file {save_path} is invalid due to unmatching {key}: "
+                        f"{f.atrs[key]} != {val}.Please recompute embeddings in a new file."
                     )
                     if wrong_file_callback is not None:
                         save_path = wrong_file_callback(save_path)