Update model combination tooltips and adapt other GUIs (#964)

anwai98 · web-flow · commit 453a8057183c · 2025-03-28T11:55:33.000+01:00
Update model combination tooltips and adapt image series annotator model and training GUI
diff --git a/micro_sam/napari.yaml b/micro_sam/napari.yaml
@@ -1,5 +1,5 @@
 name: micro-sam
-display_name: SegmentAnything for Microscopy
+display_name: Segment Anything for Microscopy
 # see https://napari.org/stable/plugins/manifest.html for valid categories
 categories: ["Segmentation", "Annotation"]
 contributions:
diff --git a/micro_sam/sam_annotator/_tooltips.py b/micro_sam/sam_annotator/_tooltips.py
@@ -8,7 +8,8 @@
     "embeddings_save_path": "Select path to save or load the computed image embeddings.",
     "halo": "Enter overlap values for computing tiled embeddings. Enter only x-value for quadratic size.\n Only active when tiling is used.",  # noqa
     "image": "Select the napari image layer.",
-    "model": "Select the segment anything model.",
+    "model_family": "Select the segment anything model family.",
+    "model_size": "Select the image encoder size of the segment anything model.",
     "automatic_segmentation_mode": "Select the automatic segmentation mode.",
     "run_button": "Compute embeddings or load embeddings if embedding_save_path is specified.",
     "tiling": "Enter tile size for computing tiled embeddings. Enter only x-value for quadratic size or both for non-quadratic.",  # noqa
diff --git a/micro_sam/sam_annotator/_widgets.py b/micro_sam/sam_annotator/_widgets.py
@@ -232,6 +232,127 @@ def _get_file_path(self, name, textbox, tooltip=None):
             # Handle the case where the selected path is not a file
             print("Invalid file selected. Please try again.")
 
+    def _get_model_size_options(self):
+        # We store the actual model names mapped to UI labels.
+        self.model_size_mapping = {}
+        if self.model_family == "Natural Images (SAM)":
+            self.model_size_options = list(self._model_size_map .values())
+            self.model_size_mapping = {self._model_size_map[k]: f"vit_{k}" for k in self._model_size_map.keys()}
+        else:
+            model_suffix = self.supported_dropdown_maps[self.model_family]
+            self.model_size_options = []
+
+            for option in self.model_options:
+                if option.endswith(model_suffix):
+                    # Extract model size character on-the-fly.
+                    key = next((k for k in self._model_size_map .keys() if f"vit_{k}" in option), None)
+                    if key:
+                        size_label = self._model_size_map[key]
+                        self.model_size_options.append(size_label)
+                        self.model_size_mapping[size_label] = option  # Store the actual model name.
+
+        # We ensure an assorted order of model sizes ('tiny' to 'huge')
+        self.model_size_options.sort(key=lambda x: ["tiny", "base", "large", "huge"].index(x))
+
+    def _update_model_type(self):
+        # Get currently selected model size (before clearing dropdown)
+        current_selection = self.model_size_dropdown.currentText()
+        self._get_model_size_options()  # Update model size options dynamically
+
+        # NOTE: We need to prevent recursive updates for this step temporarily.
+        self.model_size_dropdown.blockSignals(True)
+
+        # Let's clear and recreate the dropdown.
+        self.model_size_dropdown.clear()
+        self.model_size_dropdown.addItems(self.model_size_options)
+
+        # We restore the previous selection, if still valid.
+        if current_selection in self.model_size_options:
+            self.model_size = current_selection
+        else:
+            if self.model_size_options:  # Default to the first available model size
+                self.model_size = self.model_size_options[0]
+
+        # Let's map the selection to the correct model type (eg. "tiny" -> "vit_t")
+        size_key = next(
+            (k for k, v in self._model_size_map.items() if v == self.model_size), "b"
+        )
+        self.model_type = f"vit_{size_key}" + self.supported_dropdown_maps[self.model_family]
+
+        self.model_size_dropdown.setCurrentText(self.model_size)  # Apply the selected text to the dropdown
+
+        # We force a refresh for UI here.
+        self.model_size_dropdown.update()
+
+        # NOTE: And finally, we should re-enable signals again.
+        self.model_size_dropdown.blockSignals(False)
+
+    def _create_model_section(self, default_model: str = util._DEFAULT_MODEL, create_layout: bool = True):
+
+        # Create a list of support dropdown values and correspond them to suffixes.
+        self.supported_dropdown_maps = {
+            "Natural Images (SAM)": "",
+            "Light Microscopy": "_lm",
+            "Electron Microscopy": "_em_organelles",
+            "Medical Imaging": "_medical_imaging",
+            "Histopathology": "_histopathology",
+        }
+
+        # NOTE: The available options for all are either 'tiny', 'base', 'large' or 'huge'.
+        self._model_size_map = {"t": "tiny", "b": "base", "l": "large", "h": "huge"}
+
+        self._default_model_choice = default_model
+        # Let's set the literally default model choice depending on 'micro-sam'.
+        self.model_family = {v: k for k, v in self.supported_dropdown_maps.items()}[self._default_model_choice[5:]]
+
+        kwargs = {}
+        if create_layout:
+            layout = QtWidgets.QVBoxLayout()
+            kwargs["layout"] = layout
+
+        # NOTE: We stick to the base variant for each model family.
+        # i.e. 'Natural Images (SAM)', 'Light Microscopy', 'Electron Microscopy', 'Medical_Imaging', 'Histopathology'.
+        self.model_family_dropdown, layout = self._add_choice_param(
+            "model_family", self.model_family, list(self.supported_dropdown_maps.keys()),
+            title="Model:", tooltip=get_tooltip("embedding", "model_family"), **kwargs,
+        )
+        self.model_family_dropdown.currentTextChanged.connect(self._update_model_type)
+        return layout
+
+    def _create_model_size_section(self):
+
+        # Create UI for the model size.
+        # This would combine with the chosen 'self.model_family' and depend on 'self._default_model_choice'.
+        self.model_size = self._model_size_map[self._default_model_choice[4]]
+
+        # Get all model options.
+        self.model_options = list(util.models().urls.keys())
+        # Filter out the decoders from the model list.
+        self.model_options = [model for model in self.model_options if not model.endswith("decoder")]
+
+        # Now, we get the available sizes per model family.
+        self._get_model_size_options()
+
+        self.model_size_dropdown, layout = self._add_choice_param(
+            "model_size", self.model_size, self.model_size_options,
+            title="model size:", tooltip=get_tooltip("embedding", "model_size"),
+        )
+        self.model_size_dropdown.currentTextChanged.connect(self._update_model_type)
+        return layout
+
+    def _validate_model_type_and_custom_weights(self):
+        # Let's get all model combination stuff into the desired `model_type` structure.
+        self.model_type = "vit_" + self.model_size[0] + self.supported_dropdown_maps[self.model_family]
+
+        # For 'custom_weights', we remove the displayed text on top of the drop-down menu.
+        if self.custom_weights:
+            # NOTE: We prevent recursive updates for this step temporarily.
+            self.model_family_dropdown.blockSignals(True)
+            self.model_family_dropdown.setCurrentIndex(-1)  # This removes the displayed text.
+            self.model_family_dropdown.update()
+            # NOTE: And re-enable signals again.
+            self.model_family_dropdown.blockSignals(False)
+
 
 # Custom signals for managing progress updates.
 class PBarSignals(QObject):
@@ -1016,7 +1137,7 @@ def __init__(self, parent=None):
         # Section 1: Image and Model.
         section1_layout = QtWidgets.QHBoxLayout()
         section1_layout.addLayout(self._create_image_section())
-        section1_layout.addLayout(self._create_model_section())
+        section1_layout.addLayout(self._create_model_section())  # Creates the model family widget section.
         self.layout().addLayout(section1_layout)
 
         # Section 2: Settings (collapsible).
@@ -1103,116 +1224,15 @@ def _update_model(self, state):
         if "segment_nd" in state.widgets:
             vutil._sync_ndsegment_widget(state.widgets["segment_nd"], _model_type, self.custom_weights)
 
-    def _update_model_type(self):
-        # Get currently selected model size (before clearing dropdown)
-        current_selection = self.model_size_dropdown.currentText()
-        self._get_model_size_options()  # Update model size options dynamically
-
-        # NOTE: We need to prevent recursive updates for this step temporarily.
-        self.model_size_dropdown.blockSignals(True)
-
-        # Let's clear and recreate the dropdown.
-        self.model_size_dropdown.clear()
-        self.model_size_dropdown.addItems(self.model_size_options)
-
-        # We restore the previous selection, if still valid.
-        if current_selection in self.model_size_options:
-            self.model_size = current_selection
-        else:
-            if self.model_size_options:  # Default to the first available model size
-                self.model_size = self.model_size_options[0]
-
-        # Let's map the selection to the correct model type (eg. "tiny" -> "vit_t")
-        size_key = next(
-            (k for k, v in self._model_size_map.items() if v == self.model_size), "b"
-        )
-        self.model_type = f"vit_{size_key}" + self.supported_dropdown_maps[self.model_family]
-
-        self.model_size_dropdown.setCurrentText(self.model_size)  # Apply the selected text to the dropdown
-
-        # We force a refresh for UI here.
-        self.model_size_dropdown.update()
-
-        # NOTE: And finally, we should re-enable signals again.
-        self.model_size_dropdown.blockSignals(False)
-
-    def _get_model_size_options(self):
-        # We store the actual model names mapped to UI labels.
-        self.model_size_mapping = {}
-        if self.model_family == "Natural Images (SAM)":
-            self.model_size_options = list(self._model_size_map .values())
-            self.model_size_mapping = {self._model_size_map[k]: f"vit_{k}" for k in self._model_size_map.keys()}
-        else:
-            model_suffix = self.supported_dropdown_maps[self.model_family]
-            self.model_size_options = []
-
-            for option in self.model_options:
-                if option.endswith(model_suffix):
-                    # Extract model size character on-the-fly.
-                    key = next((k for k in self._model_size_map .keys() if f"vit_{k}" in option), None)
-                    if key:
-                        size_label = self._model_size_map[key]
-                        self.model_size_options.append(size_label)
-                        self.model_size_mapping[size_label] = option  # Store the actual model name.
-
-        # We ensure an assorted order of model sizes ('tiny' to 'huge')
-        self.model_size_options.sort(key=lambda x: ["tiny", "base", "large", "huge"].index(x))
-
-    def _create_model_section(self):
-        # Create a list of support dropdown values and correspond them to suffixes.
-        self.supported_dropdown_maps = {
-            "Natural Images (SAM)": "",
-            "Light Microscopy": "_lm",
-            "Electron Microscopy": "_em_organelles",
-            "Medical Imaging": "_medical_imaging",
-            "Histopathology": "_histopathology",
-        }
-
-        # NOTE: The available options for all are either 'tiny', 'base', 'large' or 'huge'.
-        self._model_size_map = {"t": "tiny", "b": "base", "l": "large", "h": "huge"}
-
-        self._default_model_choice = util._DEFAULT_MODEL
-        # Let's set the literally default model choice depending on 'micro-sam'.
-        self.model_family = {v: k for k, v in self.supported_dropdown_maps.items()}[self._default_model_choice[5:]]
-
-        layout = QtWidgets.QVBoxLayout()
-
-        # NOTE: We stick to the base variant for each model family.
-        # i.e. 'Natural Images (SAM)', 'Light Microscopy', 'Electron Microscopy', 'Medical_Imaging', 'Histopathology'.
-        self.model_family_dropdown, layout = self._add_choice_param(
-            "model_family", self.model_family, list(self.supported_dropdown_maps.keys()),
-            title="Model:", layout=layout, tooltip=get_tooltip("embedding", "model")
-        )
-        self.model_family_dropdown.currentTextChanged.connect(self._update_model_type)
-        return layout
-
     def _create_settings_widget(self):
         setting_values = QtWidgets.QWidget()
         setting_values.setToolTip(get_tooltip("embedding", "settings"))
         setting_values.setLayout(QtWidgets.QVBoxLayout())
 
-        # Create UI for the model size.
-        # This would combine with the chosen 'self.model_family' and depend on 'self._default_model_choice'.
-        self.model_size = self._model_size_map[self._default_model_choice[4]]
-
-        # Get all model options.
-        self.model_options = list(util.models().urls.keys())
-        # Filter out the decoders from the model list.
-        self.model_options = [model for model in self.model_options if not model.endswith("decoder")]
-
-        # Now, we get the available sizes per model family.
-        self._get_model_size_options()
-
-        self.model_size_dropdown, layout = self._add_choice_param(
-            "model_size", self.model_size, self.model_size_options,
-            title="model size:", tooltip=get_tooltip("embedding", "model"),
-        )
-        self.model_size_dropdown.currentTextChanged.connect(self._update_model_type)
+        # Add the model size widget section.
+        layout = self._create_model_size_section()
         setting_values.layout().addLayout(layout)
 
-        # Now that all parameters in place, let's get them all into one `model_type`.
-        self.model_type = "vit_" + self.model_size[0] + self.supported_dropdown_maps[self.model_family]
-
         # Create UI for the device.
         self.device = "auto"
         device_options = ["auto"] + util._available_devices()
@@ -1356,19 +1376,12 @@ def _validate_existing_embeddings(self, state):
             return _generate_message(val_results["message_type"], val_results["message"])
 
     def __call__(self, skip_validate=False):
+        self._validate_model_type_and_custom_weights()
+
         # Validate user inputs.
         if not skip_validate and self._validate_inputs():
             return
 
-        # For 'custom_weights', we remove the displayed text on top of the drop-down menu.
-        if self.custom_weights:
-            # NOTE: We prevent recursive updates for this step temporarily.
-            self.model_family_dropdown.blockSignals(True)
-            self.model_family_dropdown.setCurrentIndex(-1)  # This removes the displayed text.
-            self.model_family_dropdown.update()
-            # NOTE: And re-enable signals again.
-            self.model_family_dropdown.blockSignals(False)
-
         # Get the image.
         image = self.image_selection.get_value()
 
diff --git a/micro_sam/sam_annotator/image_series_annotator.py b/micro_sam/sam_annotator/image_series_annotator.py
@@ -370,7 +370,6 @@ def __init__(self, viewer: napari.Viewer, parent=None):
         self.run_button.clicked.connect(self.__call__)
         self.layout().addWidget(self.run_button)
 
-    # model_type: str = util._DEFAULT_MODEL,
     def _create_options(self):
         self.folder = None
         _, layout = self._add_path_param(
@@ -388,19 +387,18 @@ def _create_options(self):
         )
         self.layout().addLayout(layout)
 
-        self.model_type = util._DEFAULT_MODEL
-        model_options = list(util.models().urls.keys())
-        model_options = [model for model in model_options if not model.endswith("decoder")]
-        _, layout = self._add_choice_param(
-            "model_type", self.model_type, model_options, title="Model:",
-            tooltip=get_tooltip("embedding", "model")
-        )
+        # Add the model family widget section.
+        layout = self._create_model_section(create_layout=False)
         self.layout().addLayout(layout)
 
     def _create_settings(self):
         setting_values = QtWidgets.QWidget()
         setting_values.setLayout(QtWidgets.QVBoxLayout())
 
+        # Add the model size widget section.
+        layout = self._create_model_size_section()
+        setting_values.layout().addLayout(layout)
+
         self.pattern = "*"
         _, layout = self._add_string_param(
             "pattern", self.pattern, tooltip=get_tooltip("image_series_annotator", "pattern")
@@ -463,12 +461,16 @@ def _validate_inputs(self):
         return False
 
     def __call__(self, skip_validate=False):
+        self._validate_model_type_and_custom_weights()
+
         if not skip_validate and self._validate_inputs():
             return
         tile_shape, halo = widgets._process_tiling_inputs(self.tile_x, self.tile_y, self.halo_x, self.halo_y)
 
         image_folder_annotator(
-            self.folder, self.output_folder, self.pattern,
+            input_folder=self.folder,
+            output_folder=self.output_folder,
+            pattern=self.pattern,
             model_type=self.model_type,
             embedding_path=self.embeddings_save_path,
             tile_shape=tile_shape, halo=halo, checkpoint_path=self.custom_weights,
diff --git a/micro_sam/sam_annotator/training_ui.py b/micro_sam/sam_annotator/training_ui.py
diff --git a/test/test_sam_annotator/test_image_series_annotator.py b/test/test_sam_annotator/test_image_series_annotator.py
