Minor fix to livecell eval. scripts (#316)

* Fix evaluation script

* Update precomputing embeddings for val + other minor fixes

Author: anwai98

finetuning/livecell/evaluation/iterative_prompting.py

@@ -66,7 +66,7 @@ def main():
     predictor = get_model(model_type=args.model, ckpt=args.checkpoint)
 
     prediction_root = run_interactive_prompting(args.experiment_folder, predictor, start_with_box_prompt)
-    evaluate_interactive_prompting(prediction_root, start_with_box_prompt)
+    evaluate_interactive_prompting(prediction_root, start_with_box_prompt, args.experiment_folder)
 
 
 if __name__ == "__main__":

finetuning/livecell/evaluation/precompute_embeddings.py

@@ -2,23 +2,26 @@
 import os
 
 from micro_sam.evaluation import precompute_all_embeddings
-from util import get_paths, get_model, get_experiment_folder
+from util import get_paths, get_model
 
 
 def main():
     parser = argparse.ArgumentParser()
-    parser.add_argument("-n", "--name", required=True)
-    parser.add_argument("-m", "--model")
-    parser.add_argument("-c", "--checkpoint")
-
+    parser.add_argument("-m", "--model", type=str, required=True)
+    parser.add_argument("-c", "--checkpoint", type=str, required=True)
+    parser.add_argument("-e", "--experiment_folder", type=str, required=True)
     args = parser.parse_args()
-    name = args.name
 
-    image_paths, _ = get_paths()
-    predictor = get_model(name, model_type=args.model, ckpt=args.checkpoint)
-    exp_folder = get_experiment_folder(name)
-    embedding_dir = os.path.join(exp_folder, "embeddings")
+    predictor = get_model(model_type=args.model, ckpt=args.checkpoint)
+    embedding_dir = os.path.join(args.experiment_folder, "embeddings")
     os.makedirs(embedding_dir, exist_ok=True)
+
+    # getting the embeddings for the test set
+    image_paths, _ = get_paths("test")
+    precompute_all_embeddings(predictor, image_paths, embedding_dir)
+
+    # getting the embeddings for the val set
+    image_paths, _ = get_paths("val")
     precompute_all_embeddings(predictor, image_paths, embedding_dir)
 
 

finetuning/livecell/evaluation/precompute_embeddings.sbatch

@@ -9,4 +9,6 @@
 
 source ~/.bashrc
 micromamba activate main
-python precompute_embeddings.py $@ 
+python precompute_embeddings.py -c /scratch/usr/nimanwai/micro-sam/checkpoints/vit_b/lm_generalist_sam/best.pt \
+                                -m vit_b \
+                                -e /scratch/projects/nim00007/sam/experiments/new_models/generalists/livecell/vit_b/

finetuning/livecell/evaluation/submit_evaluation.py

@@ -5,7 +5,7 @@
 from datetime import datetime
 
 
-def write_batch_script(env_name, out_path, inference_setup, checkpoint, model_type, experiment_folder):
+def write_batch_script(env_name, out_path, inference_setup, checkpoint, model_type, experiment_folder, delay=True):
     """Writing scripts with different fold-trainings for micro-sam evaluation
     """
     batch_script = f"""#!/bin/bash
@@ -17,24 +17,40 @@ def write_batch_script(env_name, out_path, inference_setup, checkpoint, model_ty
 #SBATCH -A nim00007
 #SBATCH --job-name={inference_setup}
 
-source ~/.bashrc 
-mamba activate {env_name}
-python {inference_setup}.py """
+source ~/.bashrc
+mamba activate {env_name} \n"""
+
+    if delay:
+        batch_script += "sleep 10m \n"
+
+    # python script
+    python_script = f"python {inference_setup}.py "
 
     _op = out_path[:-3] + f"_{inference_setup}.sh"
 
     # add the finetuned checkpoint
-    batch_script += f"-c {checkpoint} "
+    python_script += f"-c {checkpoint} "
 
     # name of the model configuration
-    batch_script += f"-m {model_type} "
+    python_script += f"-m {model_type} "
 
     # experiment folder
-    batch_script += f"-e {experiment_folder} "
+    python_script += f"-e {experiment_folder} "
+
+    # let's add the python script to the bash script
+    batch_script += python_script
 
     with open(_op, "w") as f:
         f.write(batch_script)
 
+    # we run the first prompt for iterative once starting with point, and then starting with box (below)
+    if inference_setup == "iterative_prompting":
+        batch_script += "--box "
+
+        new_path = out_path[:-3] + f"_{inference_setup}_box.sh"
+        with open(new_path, "w") as f:
+            f.write(batch_script)
+
 
 def get_batch_script_names(tmp_folder):
     tmp_folder = os.path.expanduser(tmp_folder)
@@ -56,11 +72,11 @@ def submit_slurm():
 
     # parameters to run the inference scripts
     env_name = "sam"
-    checkpoint = "/scratch/usr/nimanwai/micro-sam/checkpoints/vit_h/livecell_sam/best.pt"
-    model_type = "vit_h"
-    experiment_folder = "/scratch/projects/nim00007/sam/experiments/new_models/specialists/livecell/vit_h/"
+    model_type = "vit_b"
+    checkpoint = f"/scratch/usr/nimanwai/micro-sam/checkpoints/{model_type}/lm_generalist_sam/best.pt"
+    experiment_folder = f"/scratch/projects/nim00007/sam/experiments/new_models/generalists/livecell/{model_type}/"
 
-    all_setups = ["evaluate_amg", "evaluate_instance_segmentation", "iterative_prompting"]
+    all_setups = ["precompute_embeddings", "evaluate_amg", "evaluate_instance_segmentation", "iterative_prompting"]
     for current_setup in all_setups:
         write_batch_script(
             env_name=env_name,
@@ -69,6 +85,7 @@ def submit_slurm():
             checkpoint=checkpoint,
             model_type=model_type,
             experiment_folder=experiment_folder,
+            delay=False if current_setup == "precompute_embeddings" else True
             )
 
     for my_script in glob(tmp_folder + "/*"):

finetuning/livecell/evaluation/util.py

@@ -19,8 +19,8 @@
 }
 
 
-def get_paths():
-    return _get_livecell_paths(DATA_ROOT)
+def get_paths(split="test"):
+    return _get_livecell_paths(DATA_ROOT, split=split)
 
 
 def get_checkpoint(name):