Fix disabled progressbar

micro_sam/models/vfm.py

@@ -31,11 +31,6 @@
 import torch
 import torch.nn.functional as F
 
-try:
-    from napari.utils import progress as tqdm
-except ImportError:
-    from tqdm import tqdm
-
 from bioimage_cpp.utils import Blocking
 
 from .. import util
@@ -468,7 +463,7 @@ def _handle_pbar(verbose, pbar_init, pbar_update):
 
     def init(total, description):
         nonlocal pbar
-        pbar = tqdm(total=total, desc=description, disable=not verbose)
+        pbar = util._progress_bar(total=total, desc=description, disable=not verbose)
 
     def update(n):
         if pbar is not None:

micro_sam/object_classification.py

@@ -207,7 +207,7 @@ def _compute_embedding_object_features(
         embeddings = image_embeddings["features"].squeeze()
         # AnyUp is the slow part, so show a progress bar (which also drives napari's activity dots)
         # for the single upsampling call.
-        anyup_pbar = tqdm(
+        anyup_pbar = util._progress_bar(
             total=1, desc="Upsampling with AnyUp", disable=(not verbose) or (upsampler is None)
         )
         result = _compute_object_features_impl(
@@ -239,7 +239,7 @@ def _compute_embedding_object_features(
 
     # With AnyUp, label the bar accordingly since the upsampling is the slow part.
     desc = "Upsampling with AnyUp" if upsampler is not None else "Compute object features"
-    for seg, embeds, image_region in tqdm(
+    for seg, embeds, image_region in util._progress_bar(
         seg_embed_generator(), total=n_gen, disable=not verbose, desc=desc
     ):
         # Compute this seg ids and features.

micro_sam/pixel_classification.py

@@ -266,14 +266,14 @@ def compute_pixel_features(
             ).number_of_blocks
         else:
             n_blocks = 1
-        anyup_pbar = tqdm(total=n_blocks * depth, desc="Upsampling with AnyUp", disable=not verbose)
+        anyup_pbar = util._progress_bar(total=n_blocks * depth, desc="Upsampling with AnyUp", disable=not verbose)
         pbar_update = anyup_pbar.update
     slice_disable = (not verbose) or (upsampler is not None)
 
     if is_3d:
         image_hw = (image_shape[1], image_shape[2])
         planes = []
-        for z in tqdm(range(depth), total=depth, disable=slice_disable, desc="Compute pixel features"):
+        for z in util._progress_bar(range(depth), total=depth, disable=slice_disable, desc="Compute pixel features"):
             if is_tiled:
                 plane, grid = _compute_tiled_feature_image(
                     features, image_hw, max_grid_size, z=z, image=image, upsampler=upsampler, pbar_update=pbar_update

micro_sam/precompute_state.py

@@ -14,11 +14,6 @@
 
 import torch
 
-try:
-    from napari.utils import progress as tqdm
-except ImportError:
-    from tqdm import tqdm
-
 from . import util
 from .util import AutoSegBase
 
@@ -769,7 +764,7 @@ def precompute_state(
         output_paths = [os.path.join(output_path, os.path.basename(f)) for f in input_files]
 
     predictor, current_ndim = None, None
-    for input_file, out_path in tqdm(
+    for input_file, out_path in util._progress_bar(
         zip(input_files, output_paths), total=len(input_files), desc="Precompute embeddings", disable=single
     ):
         image_data = input_file if isinstance(input_file, np.ndarray) else util.load_image_data(input_file, key)

micro_sam/util.py

@@ -46,10 +46,11 @@
     from segment_anything import sam_model_registry, SamPredictor  # noqa
     VIT_T_SUPPORT = False
 
+from tqdm import tqdm as _std_tqdm
 try:
     from napari.utils import progress as tqdm
 except ImportError:
-    from tqdm import tqdm
+    tqdm = _std_tqdm
 
 # This is the default model used in micro_sam
 # Currently it is set to vit_b_lm
@@ -725,6 +726,14 @@ def _write_embedding_signature(f, input_, predictor, tile_shape, halo, input_siz
         f.attrs[key] = val
 
 
+def _progress_bar(*args, disable=False, **kwargs):
+    # A disabled napari progress bar is never removed from napari's activity dock: its 'close' returns
+    # early when the bar is disabled. So we use a plain tqdm progress bar when the bar is disabled.
+    if disable:
+        return _std_tqdm(*args, disable=True, **kwargs)
+    return tqdm(*args, **kwargs)
+
+
 # Helper function for optional external progress bars.
 def handle_pbar(verbose, pbar_init, pbar_update):
     """@private"""

micro_sam/v1/inference.py

@@ -344,7 +344,7 @@ def _stitch_segmentation(masks, tile_ids, tiling, halo, output_shape, verbose=Fa
     assert len(masks) == len(tile_ids), f"{len(masks)}, {len(tile_ids)}"
     segmentation = np.zeros(output_shape, dtype="uint32")
 
-    for tile_id, this_seg in tqdm(zip(tile_ids, masks), desc="Stitch tiles", disable=not verbose):
+    for tile_id, this_seg in util._progress_bar(zip(tile_ids, masks), desc="Stitch tiles", disable=not verbose):
         tile = tiling.getBlockWithHalo(tile_id, list(halo)).outerBlock
         bb = tuple(slice(begin, end) for begin, end in zip(tile.begin, tile.end))
         if tile_id == 0:

micro_sam/v1/multi_dimensional_segmentation.py

@@ -30,11 +30,6 @@
 
 import torch
 
-try:
-    from napari.utils import progress as tqdm
-except ImportError:
-    from tqdm import tqdm
-
 import elf.segmentation as seg_utils
 import elf.tracking.tracking_utils as track_utils
 from elf.tracking.motile_tracking import recolor_segmentation
@@ -408,7 +403,7 @@ def _segment_slices(
     offset = 0
     segmentation = np.zeros(data.shape, dtype="uint32")
 
-    for i in tqdm(range(segmentation.shape[0]), desc="Segment slices", disable=not verbose):
+    for i in util._progress_bar(range(segmentation.shape[0]), desc="Segment slices", disable=not verbose):
         segmentor.initialize(data[i], image_embeddings=image_embeddings, verbose=False, i=i)
         seg = segmentor.generate(**kwargs)