Update generalist training and fix some issue in embedding generation

Author: constantinpape

finetuning/generalists/lm/evaluate_generalist.py

@@ -0,0 +1,18 @@
+from util import evaluate_checkpoint_for_datasets
+
+
+# TODO extend this to run the full evaluation protocol for a generalist.
+
+checkpoint = "/scratch-grete/projects/nim00007/sam/LM/generalist/vit_b/epoch-30.pt"
+root = "/scratch-grete/projects/nim00007/sam/experiments/generalists/lm/test"
+datasets = ["covid-if"]
+
+evaluate_checkpoint_for_datasets(
+    checkpoint=checkpoint,
+    model_type="vit_b",
+    experiment_root=root,
+    datasets=datasets,
+    run_default_evaluation=True,
+    run_amg=True,
+    max_num_val_images=10,
+)

finetuning/generalists/lm/evaluate_training_evolution.py

@@ -1,11 +1,9 @@
 import argparse
 import os
-import warnings
 from glob import glob
 
 import pandas as pd
-from micro_sam.util import get_custom_sam_model
-from util import evaluate_checkpoint_for_datasets
+from util import evaluate_checkpoint_for_datasets, get_generalist_predictor
 
 CHECKPOINT_ROOT = "/scratch-grete/projects/nim00007/sam/LM/generalist"
 EXPERIMENT_ROOT = "/scratch-grete/projects/nim00007/sam/experiments/generalists/lm"
@@ -26,9 +24,7 @@ def evaluate_training_evolution(model_type):
     epochs, results = [], []
     for checkpoint in checkpoints:
 
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            predictor, state = get_custom_sam_model(checkpoint, model_type=model_type, return_state=True)
+        predictor, state = get_generalist_predictor(checkpoint, model_type, return_state=True)
         epoch = state["epoch"] + 1
 
         if epoch in epochs:

finetuning/generalists/lm/util.py

@@ -1,10 +1,12 @@
+import json
 import os
+import warnings
 from glob import glob
 from pathlib import Path
 
 import pandas as pd
 from micro_sam.evaluation import (
-    inference, evaluation,
+    automatic_mask_generation, inference, evaluation,
     default_experiment_settings, get_experiment_setting_name
 )
 
@@ -22,35 +24,46 @@
 )
 
 
-def get_data_paths(dataset, split):
-    image_paths = sorted(glob(os.path.join(DATA_ROOT, dataset, split, "image_*.tif")))
+def get_generalist_predictor(checkpoint, model_type, return_state=False):
+    with warnings.catch_warnings():
+        warnings.simplefilter("ignore")
+        return inference.get_predictor(
+            checkpoint, model_type=model_type, return_state=return_state, is_custom_model=True
+        )
+
+
+def get_data_paths(dataset, split, max_num_images=None):
+    image_pattern = os.path.join(DATA_ROOT, dataset, split, "image_*.tif")
+    image_paths = sorted(glob(image_pattern))
     gt_paths = sorted(glob(os.path.join(DATA_ROOT, dataset, split, "labels_*.tif")))
     assert len(image_paths) == len(gt_paths)
-    assert len(image_paths) > 0
+    assert len(image_paths) > 0, image_pattern
+    if max_num_images is not None:
+        image_paths, gt_paths = image_paths[:max_num_images], gt_paths[:max_num_images]
     return image_paths, gt_paths
 
 
 def evaluate_checkpoint_for_dataset(
     checkpoint, model_type, dataset, experiment_folder,
     run_default_evaluation, run_amg, predictor=None,
+    max_num_val_images=None,
 ):
-    """Evaluate a generalist checkpoint for a given dataset
+    """Evaluate a generalist checkpoint for a given dataset.
     """
     assert run_default_evaluation or run_amg
 
     prompt_dir = os.path.join(PROMPT_ROOT, dataset)
 
     if predictor is None:
-        predictor = inference.get_predictor(checkpoint, model_type)
+        predictor = get_generalist_predictor(checkpoint, model_type)
     test_image_paths, test_gt_paths = get_data_paths(dataset, "test")
 
+    embedding_dir = os.path.join(experiment_folder, "test", "embeddings")
+    os.makedirs(embedding_dir, exist_ok=True)
+    result_dir = os.path.join(experiment_folder, "results")
+
     results = []
     if run_default_evaluation:
-        embedding_dir = os.path.join(experiment_folder, "test", "embeddings")
-        os.makedirs(embedding_dir, exist_ok=True)
-
-        result_dir = os.path.join(experiment_folder, "results")
-
         prompt_settings = default_experiment_settings()
         for setting in prompt_settings:
 
@@ -75,7 +88,36 @@ def evaluate_checkpoint_for_dataset(
             results.append(result)
 
     if run_amg:
-        raise NotImplementedError
+        val_embedding_dir = os.path.join(experiment_folder, "val", "embeddings")
+        val_result_dir = os.path.join(experiment_folder, "val", "results")
+        os.makedirs(val_embedding_dir, exist_ok=True)
+
+        val_image_paths, val_gt_paths = get_data_paths(dataset, "val", max_num_images=max_num_val_images)
+        automatic_mask_generation.run_amg_grid_search(
+            predictor, val_image_paths, val_gt_paths, val_embedding_dir,
+            val_result_dir, verbose_gs=True,
+        )
+
+        best_iou_thresh, best_stability_thresh, _ = automatic_mask_generation.evaluate_amg_grid_search(val_result_dir)
+        best_settings = {"pred_iou_thresh": best_iou_thresh, "stability_score_thresh": best_stability_thresh}
+        gs_result_path = os.path.join(experiment_folder, "best_gs_params.json")
+        with open(gs_result_path, "w") as f:
+            json.dump(best_settings, f)
+
+        prediction_dir = os.path.join(experiment_folder, "test", "amg")
+        os.makedirs(prediction_dir, exist_ok=True)
+        automatic_mask_generation.run_amg_inference(
+            predictor, test_image_paths, embedding_dir, prediction_dir,
+            amg_generate_kwargs=best_settings,
+        )
+
+        pred_paths = sorted(glob(os.path.join(prediction_dir, "*.tif")))
+        result_path = os.path.join(result_dir, "amg.csv")
+        os.makedirs(Path(result_path).parent, exist_ok=True)
+
+        result = evaluation.run_evaluation(test_gt_paths, pred_paths, result_path)
+        result.insert(0, "setting", ["amg"])
+        results.append(result)
 
     results = pd.concat(results)
     results.insert(0, "dataset", [dataset] * results.shape[0])
@@ -85,9 +127,10 @@ def evaluate_checkpoint_for_dataset(
 def evaluate_checkpoint_for_datasets(
     checkpoint, model_type, experiment_root, datasets,
     run_default_evaluation, run_amg, predictor=None,
+    max_num_val_images=None,
 ):
     if predictor is None:
-        predictor = inference.get_predictor(checkpoint, model_type)
+        predictor = get_generalist_predictor(checkpoint, model_type)
 
     results = []
     for dataset in datasets:
@@ -97,6 +140,7 @@ def evaluate_checkpoint_for_datasets(
             None, None, dataset, experiment_folder,
             run_default_evaluation=run_default_evaluation,
             run_amg=run_amg, predictor=predictor,
+            max_num_val_images=max_num_val_images,
         )
         results.append(result)
 

micro_sam/evaluation/automatic_mask_generation.py

@@ -118,7 +118,7 @@ def run_amg_grid_search(
         image = imageio.imread(image_path)
         gt = imageio.imread(gt_path)
 
-        embedding_path = os.path.join(embedding_dir, f"{image_name[:-4]}.zarr")
+        embedding_path = os.path.join(embedding_dir, f"{os.path.splitext(image_name)[0]}.zarr")
         image_embeddings = util.precompute_image_embeddings(predictor, image, embedding_path)
         amg.initialize(image, image_embeddings)
 
@@ -166,7 +166,7 @@ def run_amg_inference(
         assert os.path.exists(image_path), image_path
         image = imageio.imread(image_path)
 
-        embedding_path = os.path.join(embedding_dir, f"{image_name[:-4]}.zarr")
+        embedding_path = os.path.join(embedding_dir, f"{os.path.splitext(image_name)[0]}.zarr")
         image_embeddings = util.precompute_image_embeddings(predictor, image, embedding_path)
 
         amg.initialize(image, image_embeddings)

micro_sam/evaluation/inference.py

@@ -178,19 +178,24 @@ def _run_inference_with_prompts_for_image(
 def get_predictor(
     checkpoint_path: Union[str, os.PathLike],
     model_type: str,
-    return_state: bool = False
+    return_state: bool = False,
+    is_custom_model: Optional[bool] = None,
 ) -> SamPredictor:
     """Get the segment anything predictor from an exported or custom checkpoint.
 
     Args:
         checkpoint_path: The checkpoint filepath.
         model_type: The type of the model, either vit_h, vit_b or vit_l.
         return_state: Whether to return the complete state of the checkpoint in addtion to the predictor.
+        is_custom_model: Whether this is a custom model or not.
     Returns:
         The segment anything predictor.
     """
-    # TODO use try-except rather than this construct, so that we don't rely on the checkpoint name
-    if checkpoint_path.split("/")[-1] == "best.pt":  # Finetuned SAM model
+    # By default we check if the model follows the torch_em checkpint naming scheme to check whether it is a
+    # custom model or not. This can be over-ridden by passing True or False for is_custom_model.
+    is_custom_model = checkpoint_path.split("/")[-1] == "best.pt" if is_custom_model is None else is_custom_model
+
+    if is_custom_model:  # Finetuned SAM model
         predictor = util.get_custom_sam_model(
             checkpoint_path=checkpoint_path, model_type=model_type, return_state=return_state
         )
@@ -217,7 +222,7 @@ def precompute_all_embeddings(
     for image_path in tqdm(image_paths, desc="Precompute embeddings"):
         image_name = os.path.basename(image_path)
         im = imageio.imread(image_path)
-        embedding_path = os.path.join(embedding_dir, f"{image_name[:-4]}.zarr")
+        embedding_path = os.path.join(embedding_dir, f"{os.path.splitext(image_name)[0]}.zarr")
         util.precompute_image_embeddings(predictor, im, embedding_path)
 
 
@@ -384,7 +389,7 @@ def run_inference_with_prompts(
         gt = imageio.imread(gt_path).astype("uint32")
         gt = relabel_sequential(gt)[0]
 
-        embedding_path = os.path.join(embedding_dir, f"{image_name[:-4]}.zarr")
+        embedding_path = os.path.join(embedding_dir, f"{os.path.splitext(image_name)[0]}.zarr")
         image_embeddings = util.precompute_image_embeddings(predictor, im, embedding_path)
         util.set_precomputed(predictor, image_embeddings)
 

micro_sam/util.py

@@ -496,8 +496,10 @@ def precompute_image_embeddings(
             if "input_size" in f.attrs:  # we have computed the embeddings already
                 # 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 is invalid due to unmatching {key}. \
-                        Please recompute embeddings in a new file.")
+                    warnings.warn(
+                        f"Embeddings file {save_path} is invalid due to unmatching {key}."
+                        "Please recompute embeddings in a new file."
+                    )
                     if wrong_file_callback is not None:
                         save_path = wrong_file_callback(save_path)
                         f = zarr.open(save_path, "a")