Allow flexible choice between metrics for grid-search (#812)

Author: anwai98

micro_sam/evaluation/instance_segmentation.py

@@ -309,9 +309,7 @@ def run_instance_segmentation_inference(
 
 
 def evaluate_instance_segmentation_grid_search(
-    result_dir: Union[str, os.PathLike],
-    grid_search_parameters: List[str],
-    criterion: str = "mSA"
+    result_dir: Union[str, os.PathLike], grid_search_parameters: List[str], criterion: str = "mSA"
 ) -> Tuple[Dict[str, Any], float]:
     """Evaluate gridsearch results.
 
@@ -324,7 +322,6 @@ def evaluate_instance_segmentation_grid_search(
         The best parameter setting.
         The evaluation score for the best setting.
     """
-
     # Load all the grid search results.
     gs_files = glob(os.path.join(result_dir, "*.csv"))
     gs_result = pd.concat([pd.read_csv(gs_file) for gs_file in gs_files])

micro_sam/evaluation/multi_dimensional_segmentation.py

@@ -4,13 +4,13 @@
 from tqdm import tqdm
 from math import floor
 from itertools import product
-from typing import Union, Tuple, Optional, List, Dict
+from typing import Union, Tuple, Optional, List, Dict, Literal
 
 import imageio.v3 as imageio
 
 import torch
 
-from elf.evaluation import mean_segmentation_accuracy
+from elf.evaluation import mean_segmentation_accuracy, dice_score
 
 from .. import util
 from ..inference import batched_inference
@@ -30,7 +30,7 @@ def default_grid_search_values_multi_dimensional_segmentation(
         iou_threshold_values: The values for `iou_threshold` used in the grid-search.
             By default values in the range from 0.5 to 0.9 with a stepsize of 0.1 will be used.
         projection_method_values: The values for `projection` method used in the grid-search.
-            By default the values `mask`, `bounding_box` and `points` are used.
+            By default the values `mask`, `points`, `box`, `points_and_mask` and `single_point` are used.
         box_extension_values: The values for `box_extension` used in the grid-search.
             By default values in the range from 0 to 0.25 with a stepsize of 0.025 will be used.
 
@@ -71,6 +71,7 @@ def segment_slices_from_ground_truth(
     verbose: bool = False,
     return_segmentation: bool = False,
     min_size: int = 0,
+    evaluation_metric: Literal["sa", "dice"] = "sa",
 ) -> Union[float, Tuple[np.ndarray, float]]:
     """Segment all objects in a volume by prompt-based segmentation in one slice per object.
 
@@ -94,6 +95,7 @@ def segment_slices_from_ground_truth(
         return_segmentation: Whether to return the segmented volume.
         min_size: The minimal size for evaluating an object in the ground-truth.
             The size is measured within the central slice.
+        evaluation_metric: The choice of supported metric to evaluate predictions.
     """
     assert volume.ndim == 3
 
@@ -116,7 +118,7 @@ def segment_slices_from_ground_truth(
         _segmentation_completed = True  # We avoid rerunning the segmentation if it is completed.
 
     skipped_label_ids = []
-    for label_id in tqdm(label_ids, desc="Segmenting per object in the volume"):
+    for label_id in tqdm(label_ids, desc="Segmenting per object in the volume", disable=not verbose):
         # Binary label volume per instance (also referred to as object)
         this_seg = (ground_truth == label_id).astype("int")
 
@@ -204,43 +206,56 @@ def segment_slices_from_ground_truth(
     else:
         curr_gt = ground_truth
 
-    msa, sa = mean_segmentation_accuracy(final_segmentation, curr_gt, return_accuracies=True)
-    results = {"mSA": msa, "SA50": sa[0], "SA75": sa[5]}
-    results = pd.DataFrame.from_dict([results])
+    if evaluation_metric == "sa":
+        msa, sa = mean_segmentation_accuracy(
+            segmentation=final_segmentation, groundtruth=curr_gt, return_accuracies=True
+        )
+        results = {"mSA": msa, "SA50": sa[0], "SA75": sa[5]}
+    elif evaluation_metric == "dice":
+        dice = dice_score(segmentation=final_segmentation, groundtruth=curr_gt)
+        results = {"Dice": dice}
+    else:
+        raise ValueError(f"'{evaluation_metric}' is not a supported evaluation metrics. Please choose 'sa' / 'dice'.")
 
     if return_segmentation:
         return results, final_segmentation
     else:
         return results
 
 
-def _get_best_parameters_from_grid_search_combinations(result_dir, best_params_path, grid_search_values):
+def _get_best_parameters_from_grid_search_combinations(
+    result_dir, best_params_path, grid_search_values, evaluation_metric,
+):
     if os.path.exists(best_params_path):
         print("The best parameters are already saved at:", best_params_path)
         return
 
-    best_kwargs, best_msa = evaluate_instance_segmentation_grid_search(result_dir, list(grid_search_values.keys()))
+    criterion = "mSA" if evaluation_metric == "sa" else "Dice"
+    best_kwargs, best_metric = evaluate_instance_segmentation_grid_search(
+        result_dir=result_dir, grid_search_parameters=list(grid_search_values.keys()), criterion=criterion,
+    )
 
     # let's save the best parameters
-    best_kwargs["mSA"] = best_msa
+    best_kwargs[criterion] = best_metric
     best_param_df = pd.DataFrame.from_dict([best_kwargs])
     best_param_df.to_csv(best_params_path)
 
     best_param_str = ", ".join(f"{k} = {v}" for k, v in best_kwargs.items())
-    print("Best grid-search result:", best_msa, "with parmeters:\n", best_param_str)
+    print("Best grid-search result:", best_metric, "with parmeters:\n", best_param_str)
 
 
 def run_multi_dimensional_segmentation_grid_search(
     volume: np.ndarray,
     ground_truth: np.ndarray,
     model_type: str,
     checkpoint_path: Union[str, os.PathLike],
-    embedding_path: Union[str, os.PathLike],
+    embedding_path: Optional[Union[str, os.PathLike]],
     result_dir: Union[str, os.PathLike],
     interactive_seg_mode: str = "box",
     verbose: bool = False,
     grid_search_values: Optional[Dict[str, List]] = None,
-    min_size: int = 0
+    min_size: int = 0,
+    evaluation_metric: Literal["sa", "dice"] = "sa",
 ):
     """Run grid search for prompt-based multi-dimensional instance segmentation.
 
@@ -250,7 +265,7 @@ def run_multi_dimensional_segmentation_grid_search(
     ```
     grid_search_values = {
         "iou_threshold": [0.5, 0.6, 0.7, 0.8, 0.9],
-        "projection": ["mask", "bounding_box", "points"],
+        "projection": ["mask", "box", "points"],
         "box_extension": [0, 0.1, 0.2, 0.3, 0.4, 0,5],
     }
     ```
@@ -264,12 +279,13 @@ def run_multi_dimensional_segmentation_grid_search(
         model_type: Choice of segment anything model.
         checkpoint_path: Path to the model checkpoint.
         embedding_path: Path to cache the computed embeddings.
-        result_path: Path to save the grid search results.
+        result_dir: Path to save the grid search results.
         interactive_seg_mode: Method for guiding prompt-based instance segmentation.
         verbose: Whether to get the trace for projected segmentations.
         grid_search_values: The grid search values for parameters of the `segment_slices_from_ground_truth` function.
         min_size: The minimal size for evaluating an object in the ground-truth.
             The size is measured within the central slice.
+        evaluation_metric: The choice of metric for evaluating predictions.
     """
     if grid_search_values is None:
         grid_search_values = default_grid_search_values_multi_dimensional_segmentation()
@@ -280,7 +296,9 @@ def run_multi_dimensional_segmentation_grid_search(
     result_path = os.path.join(result_dir, "all_grid_search_results.csv")
     best_params_path = os.path.join(result_dir, "grid_search_params_multi_dimensional_segmentation.csv")
     if os.path.exists(result_path):
-        _get_best_parameters_from_grid_search_combinations(result_dir, best_params_path, grid_search_values)
+        _get_best_parameters_from_grid_search_combinations(
+            result_dir, best_params_path, grid_search_values, evaluation_metric
+        )
         return best_params_path
 
     # Compute all combinations of grid search values.
@@ -292,7 +310,7 @@ def run_multi_dimensional_segmentation_grid_search(
     ]
 
     net_list = []
-    for gs_kwargs in tqdm(gs_combinations):
+    for gs_kwargs in tqdm(gs_combinations, desc="Run grid-search for multi-dimensional segmentation"):
         results = segment_slices_from_ground_truth(
             volume=volume,
             ground_truth=ground_truth,
@@ -303,6 +321,7 @@ def run_multi_dimensional_segmentation_grid_search(
             verbose=verbose,
             return_segmentation=False,
             min_size=min_size,
+            evaluation_metric=evaluation_metric,
             **gs_kwargs
         )
 
@@ -313,6 +332,8 @@ def run_multi_dimensional_segmentation_grid_search(
     res_df = pd.concat(net_list, ignore_index=True)
     res_df.to_csv(result_path)
 
-    _get_best_parameters_from_grid_search_combinations(result_dir, best_params_path, grid_search_values)
+    _get_best_parameters_from_grid_search_combinations(
+        result_dir, best_params_path, grid_search_values, evaluation_metric
+    )
     print("The best grid-search parameters have been computed and stored at:", best_params_path)
     return best_params_path

micro_sam/multi_dimensional_segmentation.py

@@ -139,8 +139,9 @@ def segment_range(z_start, z_stop, increment, stopping_criterion, threshold=None
             if threshold is not None:
                 iou = util.compute_iou(seg_prev, seg_z)
                 if iou < threshold:
-                    msg = f"Segmentation stopped at slice {z} due to IOU {iou} < {threshold}."
-                    print(msg)
+                    if verbose:
+                        msg = f"Segmentation stopped at slice {z} due to IOU {iou} < {threshold}."
+                        print(msg)
                     break
 
             segmentation[z] = seg_z