Add support for tiling window in evaluation scripts (#844)

Adds support for tiling window based prediciton by allowing an optional choice for tile shape and overlap shape

Author: anwai98

micro_sam/evaluation/inference.py

@@ -6,7 +6,7 @@
 import numpy as np
 from tqdm import tqdm
 from copy import deepcopy
-from typing import Any, Dict, List, Optional, Union
+from typing import Any, Dict, List, Optional, Union, Tuple
 
 import imageio.v3 as imageio
 from skimage.segmentation import relabel_sequential
@@ -20,6 +20,7 @@
 from ..instance_segmentation import (
     mask_data_to_segmentation, get_predictor_and_decoder,
     AutomaticMaskGenerator, InstanceSegmentationWithDecoder,
+    TiledAutomaticMaskGenerator, TiledInstanceSegmentationWithDecoder,
 )
 from . import instance_segmentation
 from ..prompt_generators import PointAndBoxPromptGenerator, IterativePromptGenerator
@@ -539,7 +540,8 @@ def run_amg(
     iou_thresh_values: Optional[List[float]] = None,
     stability_score_values: Optional[List[float]] = None,
     peft_kwargs: Optional[Dict] = None,
-    cache_embeddings: bool = False
+    cache_embeddings: bool = False,
+    tiling_window_params: Optional[Dict[str, Tuple[int, int]]] = None,
 ) -> str:
     """Run Segment Anything inference for multiple images using automatic mask generation (AMG).
 
@@ -554,6 +556,7 @@ def run_amg(
         stability_score_values: Optional choice of values for grid search of `stability_score` parameter.
         peft_kwargs: Keyword arguments for th PEFT wrapper class.
         cache_embeddings: Whether to cache embeddings in experiment folder.
+        tiling_window_params: The parameters to decide whether to use tiling window operation for AIS.
 
     Returns:
         Filepath where the predictions have been saved.
@@ -566,7 +569,23 @@ def run_amg(
         embedding_folder = None
 
     predictor = util.get_sam_model(model_type=model_type, checkpoint_path=checkpoint, peft_kwargs=peft_kwargs)
-    amg = AutomaticMaskGenerator(predictor)
+
+    # Get the AMG class.
+    if tiling_window_params:
+        if not isinstance(tiling_window_params, dict):
+            raise RuntimeError("The tiling window parameters are expected to be provided as a dictionary of params.")
+
+        if "tile_shape" not in tiling_window_params:
+            raise RuntimeError("'tile_shape' parameter is missing from the provided parameters.")
+
+        if "halo" not in tiling_window_params:
+            raise RuntimeError("'halo' parameter is missing from the provided parameters.")
+
+        amg_class = TiledAutomaticMaskGenerator
+    else:
+        amg_class = AutomaticMaskGenerator
+
+    amg = amg_class(predictor)
     amg_prefix = "amg"
 
     # where the predictions are saved
@@ -592,6 +611,7 @@ def run_amg(
         prediction_dir=prediction_folder,
         result_dir=gs_result_folder,
         experiment_folder=experiment_folder,
+        tiling_window_params=tiling_window_params,
     )
     return prediction_folder
 
@@ -610,6 +630,7 @@ def run_instance_segmentation_with_decoder(
     test_image_paths: List[Union[str, os.PathLike]],
     peft_kwargs: Optional[Dict] = None,
     cache_embeddings: bool = False,
+    tiling_window_params: Optional[Dict[str, Tuple[int, int]]] = None,
 ) -> str:
     """Run Segment Anything inference for multiple images using additional automatic instance segmentation (AIS).
 
@@ -622,6 +643,7 @@ def run_instance_segmentation_with_decoder(
         test_image_paths: The list of filepaths of input images for automatic instance segmentation.
         peft_kwargs: Keyword arguments for th PEFT wrapper class.
         cache_embeddings: Whether to cache embeddings in experiment folder.
+        tiling_window_params: The parameters to decide whether to use tiling window operation for AIS.
 
     Returns:
         Filepath where the predictions have been saved.
@@ -636,7 +658,23 @@ def run_instance_segmentation_with_decoder(
     predictor, decoder = get_predictor_and_decoder(
         model_type=model_type, checkpoint_path=checkpoint, peft_kwargs=peft_kwargs,
     )
-    segmenter = InstanceSegmentationWithDecoder(predictor, decoder)
+
+    # Get the AIS class.
+    if tiling_window_params:
+        if not isinstance(tiling_window_params, dict):
+            raise RuntimeError("The tiling window parameters are expected to be provided as a dictionary of params.")
+
+        if "tile_shape" not in tiling_window_params:
+            raise RuntimeError("'tile_shape' parameter is missing from the provided parameters.")
+
+        if "halo" not in tiling_window_params:
+            raise RuntimeError("'halo' parameter is missing from the provided parameters.")
+
+        ais_class = TiledInstanceSegmentationWithDecoder
+    else:
+        ais_class = InstanceSegmentationWithDecoder
+
+    segmenter = ais_class(predictor, decoder)
     seg_prefix = "instance_segmentation_with_decoder"
 
     # where the predictions are saved
@@ -650,9 +688,15 @@ def run_instance_segmentation_with_decoder(
     grid_search_values = instance_segmentation.default_grid_search_values_instance_segmentation_with_decoder()
 
     instance_segmentation.run_instance_segmentation_grid_search_and_inference(
-        segmenter, grid_search_values,
-        val_image_paths, val_gt_paths, test_image_paths,
-        embedding_dir=embedding_folder, prediction_dir=prediction_folder,
-        result_dir=gs_result_folder, experiment_folder=experiment_folder,
+        segmenter=segmenter,
+        grid_search_values=grid_search_values,
+        val_image_paths=val_image_paths,
+        val_gt_paths=val_gt_paths,
+        test_image_paths=test_image_paths,
+        embedding_dir=embedding_folder,
+        prediction_dir=prediction_folder,
+        result_dir=gs_result_folder,
+        experiment_folder=experiment_folder,
+        tiling_window_params=tiling_window_params,
     )
     return prediction_folder

micro_sam/evaluation/instance_segmentation.py

@@ -156,6 +156,7 @@ def run_instance_segmentation_grid_search(
     image_key: Optional[str] = None,
     gt_key: Optional[str] = None,
     rois: Optional[Tuple[slice, ...]] = None,
+    tiling_window_params: Optional[Dict[str, Tuple[int, int]]] = None,
 ) -> None:
     """Run grid search for automatic mask generation.
 
@@ -188,6 +189,7 @@ def run_instance_segmentation_grid_search(
         gt_key: Key for loading the ground-truth data from a more complex file format like HDF5.
             If not given a simple image format like tif is assumed.
         rois: Region of interests to resetrict the evaluation to.
+        tiling_window_params: The parameters to decide whether to use tiling window operation for AIS.
     """
     verbose_embeddings = False
 
@@ -233,11 +235,14 @@ def run_instance_segmentation_grid_search(
             assert predictor is not None
             embedding_path = os.path.join(embedding_dir, f"{os.path.splitext(image_name)[0]}.zarr")
 
+        if tiling_window_params is None:
+            tiling_window_params = {}
+
         image_embeddings = util.precompute_image_embeddings(
-            predictor, image, embedding_path, ndim=2, verbose=verbose_embeddings
+            predictor, image, embedding_path, ndim=2, verbose=verbose_embeddings, **tiling_window_params
         )
 
-        segmenter.initialize(image, image_embeddings)
+        segmenter.initialize(image, image_embeddings, **tiling_window_params)
 
         _grid_search_iteration(
             segmenter, gs_combinations, gt, image_name,
@@ -251,6 +256,7 @@ def run_instance_segmentation_inference(
     embedding_dir: Optional[Union[str, os.PathLike]],
     prediction_dir: Union[str, os.PathLike],
     generate_kwargs: Optional[Dict[str, Any]] = None,
+    tiling_window_params: Optional[Dict[str, Tuple[int, int]]] = None,
 ) -> None:
     """Run inference for automatic mask generation.
 
@@ -260,6 +266,8 @@ def run_instance_segmentation_inference(
         embedding_dir: Folder to cache the image embeddings.
         prediction_dir: Folder to save the predictions.
         generate_kwargs: The keyword arguments for the `generate` method of the segmenter.
+        tiling_window_params: The parameters to decide whether to use tiling window operation
+            for automatic segmentation.
     """
 
     verbose_embeddings = False
@@ -285,11 +293,14 @@ def run_instance_segmentation_inference(
             assert predictor is not None
             embedding_path = os.path.join(embedding_dir, f"{os.path.splitext(image_name)[0]}.zarr")
 
+        if tiling_window_params is None:
+            tiling_window_params = {}
+
         image_embeddings = util.precompute_image_embeddings(
-            predictor, image, embedding_path, ndim=2, verbose=verbose_embeddings
+            predictor, image, embedding_path, ndim=2, verbose=verbose_embeddings, **tiling_window_params
         )
 
-        segmenter.initialize(image, image_embeddings)
+        segmenter.initialize(image, image_embeddings, **tiling_window_params)
 
         masks = segmenter.generate(**generate_kwargs)
 
@@ -372,6 +383,7 @@ def run_instance_segmentation_grid_search_and_inference(
     result_dir: Union[str, os.PathLike],
     fixed_generate_kwargs: Optional[Dict[str, Any]] = None,
     verbose_gs: bool = True,
+    tiling_window_params: Optional[Dict[str, Tuple[int, int]]] = None,
 ) -> None:
     """Run grid search and inference for automatic mask generation.
 
@@ -390,11 +402,19 @@ def run_instance_segmentation_grid_search_and_inference(
         result_dir: Folder to cache the evaluation results per image.
         fixed_generate_kwargs: Fixed keyword arguments for the `generate` method of the segmenter.
         verbose_gs: Whether to run the gridsearch for individual images in a verbose mode.
+        tiling_window_params: The parameters to decide whether to use tiling window operation
+            for automatic segmentation.
     """
     run_instance_segmentation_grid_search(
-        segmenter, grid_search_values, val_image_paths, val_gt_paths,
-        result_dir=result_dir, embedding_dir=embedding_dir,
-        fixed_generate_kwargs=fixed_generate_kwargs, verbose_gs=verbose_gs,
+        segmenter=segmenter,
+        grid_search_values=grid_search_values,
+        image_paths=val_image_paths,
+        gt_paths=val_gt_paths,
+        result_dir=result_dir,
+        embedding_dir=embedding_dir,
+        fixed_generate_kwargs=fixed_generate_kwargs,
+        verbose_gs=verbose_gs,
+        tiling_window_params=tiling_window_params,
     )
 
     best_kwargs, best_msa = evaluate_instance_segmentation_grid_search(result_dir, list(grid_search_values.keys()))
@@ -408,5 +428,10 @@ def run_instance_segmentation_grid_search_and_inference(
     generate_kwargs.update(best_kwargs)
 
     run_instance_segmentation_inference(
-        segmenter, test_image_paths, embedding_dir, prediction_dir, generate_kwargs
+        segmenter=segmenter,
+        image_paths=test_image_paths,
+        embedding_dir=embedding_dir,
+        prediction_dir=prediction_dir,
+        generate_kwargs=generate_kwargs,
+        tiling_window_params=tiling_window_params,
     )