Merge branch 'dev' of https://github.com/computational-cell-analytics/micro-sam into dev

Author: constantinpape

micro_sam/sam_annotator/_widgets.py

@@ -21,7 +21,10 @@
 from superqt import QCollapsible
 from magicgui import magic_factory
 from magicgui.widgets import ComboBox, Container, create_widget
-from napari.qt.threading import thread_worker
+# We have disabled the thread workers for now because they result in a
+# massive slowdown in napari >= 0.5.
+# See also https://forum.image.sc/t/napari-thread-worker-leads-to-massive-slowdown/103786
+# from napari.qt.threading import thread_worker
 from napari.utils import progress
 
 from ._state import AnnotatorState
@@ -1088,7 +1091,7 @@ def __call__(self, skip_validate=False):
         # Set up progress bar and signals for using it within a threadworker.
         pbar, pbar_signals = _create_pbar_for_threadworker()
 
-        @thread_worker()
+        # @thread_worker()
         def compute_image_embedding():
 
             def pbar_init(total, description):
@@ -1103,10 +1106,12 @@ def pbar_init(total, description):
             )
             pbar_signals.pbar_stop.emit()
 
-        worker = compute_image_embedding()
-        worker.returned.connect(self._update_model)
-        worker.start()
-        return worker
+        compute_image_embedding()
+        self._update_model()
+        # worker = compute_image_embedding()
+        # worker.returned.connect(self._update_model)
+        # worker.start()
+        # return worker
 
 
 #
@@ -1195,7 +1200,7 @@ def _run_tracking(self):
         state = AnnotatorState()
         pbar, pbar_signals = _create_pbar_for_threadworker()
 
-        @thread_worker
+        # @thread_worker
         def tracking_impl():
             shape = state.image_shape
 
@@ -1237,15 +1242,17 @@ def update_segmentation(ret_val):
             self._viewer.layers["current_object"].data[seg == 1] = state.current_track_id
             self._viewer.layers["current_object"].refresh()
 
-        worker = tracking_impl()
-        worker.returned.connect(update_segmentation)
-        worker.start()
-        return worker
+        ret_val = tracking_impl()
+        update_segmentation(ret_val)
+        # worker = tracking_impl()
+        # worker.returned.connect(update_segmentation)
+        # worker.start()
+        # return worker
 
     def _run_volumetric_segmentation(self):
         pbar, pbar_signals = _create_pbar_for_threadworker()
 
-        @thread_worker
+        # @thread_worker
         def volumetric_segmentation_impl():
             state = AnnotatorState()
             shape = state.image_shape
@@ -1277,10 +1284,13 @@ def update_segmentation(seg):
             self._viewer.layers["current_object"].data = seg
             self._viewer.layers["current_object"].refresh()
 
-        worker = volumetric_segmentation_impl()
-        worker.returned.connect(update_segmentation)
-        worker.start()
-        return worker
+        seg = volumetric_segmentation_impl()
+        self._viewer.layers["current_object"].data = seg
+        self._viewer.layers["current_object"].refresh()
+        # worker = volumetric_segmentation_impl()
+        # worker.returned.connect(update_segmentation)
+        # worker.start()
+        # return worker
 
     def __call__(self):
         if _validate_embeddings(self._viewer):
@@ -1522,7 +1532,7 @@ def _empty_segmentation_warning(self):
     def _run_segmentation_2d(self, kwargs, i=None):
         pbar, pbar_signals = _create_pbar_for_threadworker()
 
-        @thread_worker
+        # @thread_worker
         def seg_impl():
             def pbar_init(total, description):
                 pbar_signals.pbar_total.emit(total)
@@ -1548,10 +1558,12 @@ def update_segmentation(seg):
                 self._viewer.layers["auto_segmentation"].data[i] = seg
             self._viewer.layers["auto_segmentation"].refresh()
 
-        worker = seg_impl()
-        worker.returned.connect(update_segmentation)
-        worker.start()
-        return worker
+        seg = seg_impl()
+        update_segmentation(seg)
+        # worker = seg_impl()
+        # worker.returned.connect(update_segmentation)
+        # worker.start()
+        # return worker
 
     # We refuse to run 3D segmentation with the AMG unless we have a GPU or all embeddings
     # are precomputed. Otherwise this would take too long.
@@ -1578,7 +1590,7 @@ def _run_segmentation_3d(self, kwargs):
 
         pbar, pbar_signals = _create_pbar_for_threadworker()
 
-        @thread_worker
+        # @thread_worker
         def seg_impl():
             segmentation = np.zeros_like(self._viewer.layers["auto_segmentation"].data)
             offset = 0
@@ -1617,10 +1629,12 @@ def update_segmentation(segmentation):
             self._viewer.layers["auto_segmentation"].data = segmentation
             self._viewer.layers["auto_segmentation"].refresh()
 
-        worker = seg_impl()
-        worker.returned.connect(update_segmentation)
-        worker.start()
-        return worker
+        seg = seg_impl()
+        update_segmentation(seg)
+        # worker = seg_impl()
+        # worker.returned.connect(update_segmentation)
+        # worker.start()
+        # return worker
 
     def __call__(self):
         if _validate_embeddings(self._viewer):

micro_sam/sam_annotator/training_ui.py

@@ -2,7 +2,7 @@
 import warnings
 
 from qtpy import QtWidgets
-from napari.qt.threading import thread_worker
+# from napari.qt.threading import thread_worker
 
 import torch
 from torch.utils.data import random_split
@@ -238,7 +238,7 @@ def __call__(self, skip_validate=False):
         else:
             checkpoint_path = self.checkpoint
 
-        @thread_worker()
+        # @thread_worker()
         def run_training():
             train_loader, val_loader = self._get_loaders()
             train_sam_for_configuration(
@@ -296,7 +296,9 @@ def run_training():
                 pbar_signals.pbar_stop.emit()
                 return export_checkpoint
 
-        worker = run_training()
-        worker.returned.connect(lambda path: print(f"Training has finished. The trained model is saved at {path}."))
-        worker.start()
-        return worker
+        path = run_training()
+        print(f"Training has finished. The trained model is saved at {path}.")
+        # worker = run_training()
+        # worker.returned.connect(lambda path: print(f"Training has finished. The trained model is saved at {path}."))
+        # worker.start()
+        # return worker

micro_sam/sam_annotator/util.py

@@ -119,6 +119,8 @@ def clear_annotations(viewer: napari.Viewer, clear_segmentations=True) -> None:
     viewer.layers["point_prompts"].refresh()
     if "prompts" in viewer.layers:
         # Select all prompts and then remove them.
+        # This is how it worked before napari 0.5.
+        # viewer.layers["prompts"].data = []
         viewer.layers["prompts"].selected_data = set(range(len(viewer.layers["prompts"].data)))
         viewer.layers["prompts"].remove_selected()
         viewer.layers["prompts"].refresh()