Implement progress bar

Author: constantinpape

micro_sam/sam_annotator/annotator_3d.py

@@ -3,6 +3,7 @@
 
 from magicgui import magicgui
 from napari import Viewer
+from napari.utils import progress
 
 from .. import util
 from ..segment_from_prompts import segment_from_mask, segment_from_points
@@ -21,14 +22,19 @@
 # TODO refactor
 def _segment_volume(
     seg, predictor, image_embeddings, segmented_slices,
-    stop_lower, stop_upper, iou_threshold, method
+    stop_lower, stop_upper, iou_threshold, method,
+    progress_bar=None,
 ):
     assert method in ("mask", "bounding_box")
     if method == "mask":
         use_mask, use_box = True, True
     else:
         use_mask, use_box = False, True
 
+    def _update_progress():
+        if progress_bar is not None:
+            progress_bar.update(1)
+
     # TODO refactor to utils so that it can be used by other plugins
     def segment_range(z_start, z_stop, increment, stopping_criterion, threshold=None, verbose=False):
         z = z_start + increment
@@ -50,6 +56,7 @@ def segment_range(z_start, z_stop, increment, stopping_criterion, threshold=None
                 if verbose:
                     print(f"Segment {z_start} to {z_stop}: stop at slice {z}")
                 break
+            _update_progress()
 
     z0, z1 = int(segmented_slices.min()), int(segmented_slices.max())
 
@@ -75,6 +82,7 @@ def segment_range(z_start, z_stop, increment, stopping_criterion, threshold=None
                 seg_prompt = np.logical_or(seg[z_start] == 1, seg[z_stop] == 1)
                 seg[z] = segment_from_mask(predictor, seg_prompt, image_embeddings=image_embeddings, i=z,
                                            use_mask=use_mask, use_box=use_box)
+                _update_progress()
 
             else:  # there is a range of more than 2 slices in between -> segment ranges
                 # segment from bottom
@@ -89,6 +97,7 @@ def segment_range(z_start, z_stop, increment, stopping_criterion, threshold=None
                     seg_prompt = np.logical_or(seg[z_mid - 1] == 1, seg[z_mid + 1] == 1)
                     seg[z_mid] = segment_from_mask(predictor, seg_prompt, image_embeddings=image_embeddings, i=z_mid,
                                                    use_mask=use_mask, use_box=use_box)
+                    _update_progress()
 
     return seg
 
@@ -118,16 +127,20 @@ def segment_slice_wigdet(v: Viewer):
 def segment_volume_widget(v: Viewer, iou_threshold: float = 0.8, method: str = "mask"):
     # step 1: segment all slices with prompts
     shape = v.layers["raw"].data.shape
-    seg, slices, stop_lower, stop_upper = segment_slices_with_prompts(
-        PREDICTOR, v.layers["prompts"], IMAGE_EMBEDDINGS, shape
-    )
 
-    # step 2: segment the rest of the volume based on smart prompting
-    seg = _segment_volume(
-        seg, PREDICTOR, IMAGE_EMBEDDINGS, slices,
-        stop_lower, stop_upper,
-        iou_threshold=iou_threshold, method=method,
-    )
+    with progress(total=shape[0]) as progress_bar:
+
+        seg, slices, stop_lower, stop_upper = segment_slices_with_prompts(
+            PREDICTOR, v.layers["prompts"], IMAGE_EMBEDDINGS, shape, progress_bar=progress_bar,
+        )
+
+        # step 2: segment the rest of the volume based on smart prompting
+        seg = _segment_volume(
+            seg, PREDICTOR, IMAGE_EMBEDDINGS, slices,
+            stop_lower, stop_upper,
+            iou_threshold=iou_threshold, method=method,
+            progress_bar=progress_bar,
+        )
 
     v.layers["current_object"].data = seg
     v.layers["current_object"].refresh()

micro_sam/sam_annotator/annotator_tracking.py

@@ -3,6 +3,7 @@
 
 from magicgui import magicgui
 from napari import Viewer
+from napari.utils import progress
 
 from .. import util
 from ..segment_from_prompts import segment_from_mask, segment_from_points
@@ -19,13 +20,17 @@
 
 # TODO motion model!!!
 # TODO handle divison annotations + division classifier
-def _track_from_prompts(seg, predictor, slices, image_embeddings, stop_upper, threshold, method):
+def _track_from_prompts(seg, predictor, slices, image_embeddings, stop_upper, threshold, method, progress_bar=None):
     assert method in ("mask", "bounding_box")
     if method == "mask":
         use_mask, use_box = True, True
     else:
         use_mask, use_box = False, True
 
+    def _update_progress():
+        if progress_bar is not None:
+            progress_bar.update(1)
+
     t0 = int(slices.min())
     t = t0 + 1
     while True:
@@ -36,6 +41,7 @@ def _track_from_prompts(seg, predictor, slices, image_embeddings, stop_upper, th
             seg_prev = seg[t - 1]
             seg_t = segment_from_mask(predictor, seg_prev, image_embeddings=image_embeddings, i=t,
                                       use_mask=use_mask, use_box=use_box)
+            _update_progress()
 
         if (threshold is not None) and (seg_prev is not None):
             iou = util.compute_iou(seg_prev, seg_t)
@@ -79,14 +85,18 @@ def segment_frame_wigdet(v: Viewer):
 
 @magicgui(call_button="Track Object [V]", method={"choices": ["bounding_box", "mask"]})
 def track_objet_widget(v: Viewer, iou_threshold: float = 0.8, method: str = "mask"):
-    # step 1: segment all slices with prompts
     shape = v.layers["raw"].data.shape
-    seg, slices, _, stop_upper = segment_slices_with_prompts(
-        PREDICTOR, v.layers["prompts"], IMAGE_EMBEDDINGS, shape
-    )
 
-    # step 2: track the object starting from the lowest annotated slice
-    seg = _track_from_prompts(seg, PREDICTOR, slices, IMAGE_EMBEDDINGS, stop_upper, iou_threshold, method)
+    with progress(total=shape[0]) as progress_bar:
+        # step 1: segment all slices with prompts
+        seg, slices, _, stop_upper = segment_slices_with_prompts(
+            PREDICTOR, v.layers["prompts"], IMAGE_EMBEDDINGS, shape, progress_bar=progress_bar
+        )
+
+        # step 2: track the object starting from the lowest annotated slice
+        seg = _track_from_prompts(
+            seg, PREDICTOR, slices, IMAGE_EMBEDDINGS, stop_upper, iou_threshold, method, progress_bar=progress_bar
+        )
 
     v.layers["current_track"].data = seg
     v.layers["current_track"].refresh()

micro_sam/sam_annotator/util.py

@@ -76,29 +76,37 @@ def prompt_layer_to_points(prompt_layer, i=None):
     return this_points, this_labels
 
 
-def segment_slices_with_prompts(predictor, prompt_layer, image_embeddings, shape):
+def segment_slices_with_prompts(predictor, prompt_layer, image_embeddings, shape, progress_bar=None):
     seg = np.zeros(shape, dtype="uint32")
 
     slices = np.unique(prompt_layer.data[:, 0]).astype("int")
     stop_lower, stop_upper = False, False
 
+    def _update_progress():
+        if progress_bar is not None:
+            progress_bar.update(1)
+
     for i in slices:
         prompts_i = prompt_layer_to_points(prompt_layer, i)
 
         # TODO also take into account division properties once we have this implemented in tracking
         # do we end the segmentation at the outer slices?
         if prompts_i is None:
+
             if i == slices[0]:
                 stop_lower = True
             elif i == slices[-1]:
                 stop_upper = True
             else:
                 raise RuntimeError("Stop slices can only be at the start or end")
+
             seg[i] = 0
+            _update_progress()
             continue
 
         points, labels = prompts_i
         seg_i = segment_from_points(predictor, points, labels, image_embeddings=image_embeddings, i=i)
         seg[i] = seg_i
+        _update_progress()
 
     return seg, slices, stop_lower, stop_upper

micro_sam/util.py

@@ -10,7 +10,11 @@
 from skimage.measure import regionprops
 
 from segment_anything import sam_model_registry, SamPredictor
-from tqdm import tqdm
+
+try:
+    from napari.utils import progress as tqdm
+except ImportError:
+    from tqdm import tqdm
 
 MODEL_URLS = {
     "vit_h": "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth",