Enable switching automatic tracking model and solver (#1267)

* Enable switching automatic tracking model and solver

* Fix iterative prompt generator test assertions

* Try to fix BioImage.IO model validation on macOS

Author: anwai98

micro_sam/automatic_segmentation.py

@@ -98,6 +98,8 @@ def automatic_tracking(
     return_embeddings: bool = False,
     annotate: bool = False,
     batch_size: int = 1,
+    mode: str = "greedy",
+    tracking_model: str = "general_2d",
     **generate_kwargs
 ) -> Tuple[np.ndarray, List[Dict]]:
     """Run automatic tracking for the input timeseries.
@@ -120,6 +122,9 @@ def automatic_tracking(
             By default, does not activate the annotator.
         batch_size: The batch size to compute image embeddings over tiles / z-planes.
             By default, does it sequentially, i.e. one after the other.
+        mode: The trackastra linking solver. One of 'greedy_nodiv', 'greedy' or 'ilp'.
+            'ilp' uses the motile solver. By default, set to 'greedy'.
+        tracking_model: The pretrained trackastra model to use. By default, set to 'general_2d'.
         generate_kwargs: optional keyword arguments for the generate function of the AMG, APG, or AIS class.
 
     Returns:
@@ -146,6 +151,8 @@ def automatic_tracking(
         halo=halo,
         verbose=verbose,
         batch_size=batch_size,
+        mode=mode,
+        tracking_model=tracking_model,
         return_embeddings=True,
         output_folder=output_path,
         **generate_kwargs,

micro_sam/bioimageio/model_export.py

@@ -221,7 +221,8 @@ def _check_model(model_description, input_paths, result_paths):
     # Load outputs.
     mask = np.load(result_paths["mask"])
 
-    with bioimageio.core.create_prediction_pipeline(model_description) as pp:
+    # Match the device used to generate the reference outputs.
+    with bioimageio.core.create_prediction_pipeline(model_description, devices=["cpu"]) as pp:
 
         # Check with all prompts. We only check the result for this setting,
         # because this was used to generate the test data.

micro_sam/multi_dimensional_segmentation.py

@@ -569,9 +569,9 @@ def _filter_lineages(lineages, tracking_result):
     return filtered_lineages
 
 
-def _tracking_impl(timeseries, segmentation, mode, min_time_extent, output_folder=None):
+def _tracking_impl(timeseries, segmentation, mode, min_time_extent, tracking_model="general_2d", output_folder=None):
     device = "cuda" if torch.cuda.is_available() else "cpu"
-    model = Trackastra.from_pretrained("general_2d", device=device)
+    model = Trackastra.from_pretrained(tracking_model, device=device)
     result = model.track(timeseries, segmentation, mode=mode)
     try:
         lineage_graph, _ = result
@@ -609,6 +609,8 @@ def track_across_frames(
     segmentation: np.ndarray,
     gap_closing: Optional[int] = None,
     min_time_extent: Optional[int] = None,
+    mode: str = "greedy",
+    tracking_model: str = "general_2d",
     verbose: bool = True,
     pbar_init: Optional[callable] = None,
     pbar_update: Optional[callable] = None,
@@ -626,6 +628,9 @@ def track_across_frames(
         gap_closing: If given, gaps in the segmentation are closed with a binary closing
             operation. The value is used to determine the number of iterations for the closing.
         min_time_extent: Require a minimal extent in time for the tracked objects.
+        mode: The trackastra linking solver. One of 'greedy_nodiv', 'greedy' or 'ilp'.
+            'ilp' uses the motile solver. By default, set to 'greedy'.
+        tracking_model: The pretrained trackastra model to use. By default, set to 'general_2d'.
         verbose: Verbosity flag. By default, set to 'True'.
         pbar_init: Function to initialize the progress bar.
         pbar_update: Function to update the progress bar.
@@ -650,8 +655,9 @@ def track_across_frames(
     segmentation, lineage = _tracking_impl(
         timeseries=np.asarray(timeseries),
         segmentation=segmentation,
-        mode="greedy",
+        mode=mode,
         min_time_extent=min_time_extent,
+        tracking_model=tracking_model,
         output_folder=output_folder,
     )
     return segmentation, lineage
@@ -664,6 +670,8 @@ def automatic_tracking_implementation(
     embedding_path: Optional[Union[str, os.PathLike]] = None,
     gap_closing: Optional[int] = None,
     min_time_extent: Optional[int] = None,
+    mode: str = "greedy",
+    tracking_model: str = "general_2d",
     tile_shape: Optional[Tuple[int, int]] = None,
     halo: Optional[Tuple[int, int]] = None,
     verbose: bool = True,
@@ -685,6 +693,9 @@ def automatic_tracking_implementation(
         gap_closing: If given, gaps in the segmentation are closed with a binary closing
             operation. The value is used to determine the number of iterations for the closing.
         min_time_extent: Require a minimal extent in time for the tracked objects.
+        mode: The trackastra linking solver. One of 'greedy_nodiv', 'greedy' or 'ilp'.
+            'ilp' uses the motile solver. By default, set to 'greedy'.
+        tracking_model: The pretrained trackastra model to use. By default, set to 'general_2d'.
         tile_shape: Shape of the tiles for tiled prediction. By default prediction is run without tiling.
         halo: Overlap of the tiles for tiled prediction. By default prediction is run without tiling.
         verbose: Verbosity flag. By default, set to 'True'.
@@ -715,6 +726,8 @@ def automatic_tracking_implementation(
         segmentation=segmentation,
         gap_closing=gap_closing,
         min_time_extent=min_time_extent,
+        mode=mode,
+        tracking_model=tracking_model,
         verbose=verbose,
         output_folder=output_folder,
     )

test/test_prompt_generators.py

@@ -176,9 +176,14 @@ def _deform_labels(labels):
 
                 self.assertTrue(agree.all())
 
-                # the condition only holds if we have a negative area (prediction mask where we don't have true mask)
-                if ((pred_mask - mask) > 0).sum() > 0:
-                    self.assertTrue(diff.all())
+                # Positive and negative prompts correct false-negative and false-positive regions, respectively.
+                has_false_negatives = ((mask - pred_mask) > 0).any()
+                has_false_positives = ((pred_mask - mask) > 0).any()
+
+                if has_false_negatives:
+                    self.assertTrue(diff[0])
+                if has_false_positives:
+                    self.assertTrue(diff[1])
 
 
 if __name__ == "__main__":