Merge pull request #115 from computational-cell-analytics/99_custom_models_gui_support

99_custom_models_gui_support

Author: constantinpape

micro_sam/sam_annotator/annotator.py

@@ -4,11 +4,12 @@
 import os
 import magicgui
 import numpy as np
+import zarr
 from magicgui.widgets import Container, Label, LineEdit, SpinBox, ComboBox
 from magicgui.application import use_app
 from PyQt5.QtWidgets import QFileDialog, QMessageBox
 
-from ..util import load_image_data
+from ..util import load_image_data, get_model_names
 from .annotator_2d import annotator_2d
 from .annotator_3d import annotator_3d
 from .image_series_annotator import image_folder_annotator
@@ -37,6 +38,19 @@ def file_is_hirarchical(path_s):
         return os.path.splitext(path_s)[1] in [".hdf5", ".h5", "n5", ".zarr"]
 
 
+def _set_embeddings_file_attributes(embeddings_path: str, cb_model, re_tile_x, re_tile_y):
+    f = zarr.open(embeddings_path, "a")
+    if "tile_shape" in f.attrs:
+        if f.attrs["tile_shape"] is None:
+            re_tile_x.value = 0
+            re_tile_y.value = 0
+        else:
+            re_tile_x.value = f.attrs["tile_shape"][0]
+            re_tile_y.value = f.attrs["tile_shape"][1]
+    if "model_type" in f.attrs:
+        cb_model.value = f.attrs["model_type"]
+
+
 @magicgui.magicgui(call_button="2d annotator", labels=False)
 def _on_2d():
     global config_dict
@@ -46,6 +60,11 @@ def _on_2d():
 
     le_file_key_img = LineEdit(value="*", label="File key input")
     le_file_key_segm = LineEdit(value="*", label="File key segmentation")
+    re_halo_x = SpinBox(value=0, max=10000, label="Halo x")
+    re_halo_y = SpinBox(value=0, max=10000, label="Halo y")
+    re_tile_x = SpinBox(value=0, max=10000, label="Tile x")
+    re_tile_y = SpinBox(value=0, max=10000, label="Tile y")
+    cb_model = ComboBox(value="vit_h_lm", choices=get_model_names(), label="Model Type")
 
     @magicgui.magicgui(call_button="Select image", labels=False)
     def on_select_image():
@@ -73,6 +92,7 @@ def on_select_embed():
                 show_error("Precompute embeddings file does not exist or has wrong file extension.")
                 return
             args["embedding_path"] = path
+            _set_embeddings_file_attributes(path, cb_model, re_tile_x, re_tile_y)
         except Exception as e:
             show_error(str(e))
 
@@ -93,12 +113,6 @@ def on_select_segm():
     pb_embed_sel = Container(widgets=[on_select_embed], layout="horizontal", labels=False)
     pb_seg_segm = Container(widgets=[on_select_segm], layout="horizontal", labels=False)
 
-    re_halo_x = SpinBox(value=0, max=10000, label="Halo x")
-    re_halo_y = SpinBox(value=0, max=10000, label="Halo y")
-    re_tile_x = SpinBox(value=0, max=10000, label="Tile x")
-    re_tile_y = SpinBox(value=0, max=10000, label="Tile y")
-    cb_model = ComboBox(value="vit_h", choices=["vit_h", "vit_l", "vit_b"], label="Model Type")
-
     @magicgui.magicgui(call_button="2d annotator", labels=False)
     def on_start():
         try:
@@ -125,6 +139,7 @@ def on_start():
     sub_widget = Container(widgets=[Container(widgets=[on_start], layout="horizontal", labels=False),
                                     pb_img_sel, lbl_opt, pb_embed_sel, pb_seg_segm, le_file_key_img, le_file_key_segm,
                                     cb_model, re_tile_x, re_tile_y, re_halo_x, re_halo_y])
+    sub_widget.root_native_widget.setWindowTitle("Segment Anything for Microscopy")
     main_widget.close()
     sub_widget.show()
 
@@ -137,6 +152,12 @@ def _on_3d():
     le_file_key_img = LineEdit(value="*", label="File key input")
     le_file_key_segm = LineEdit(value="*", label="File key segmentation")
 
+    re_halo_x = SpinBox(value=0, max=10000, label="Halo x")
+    re_halo_y = SpinBox(value=0, max=10000, label="Halo y")
+    re_tile_x = SpinBox(value=0, max=10000, label="Tile x")
+    re_tile_y = SpinBox(value=0, max=10000, label="Tile y")
+    cb_model = ComboBox(value="vit_h_lm", choices=get_model_names(), label="Model Type")
+
     @magicgui.magicgui(call_button="Select images", labels=False)
     def on_select_image():
         try:
@@ -175,6 +196,7 @@ def on_select_embed():
                 show_error("Precompute embeddings file does not exist or has wrong file extension.")
                 return
             args["embedding_path"] = path
+            _set_embeddings_file_attributes(path, cb_model, re_tile_x, re_tile_y)
         except Exception as e:
             show_error(str(e))
 
@@ -208,12 +230,6 @@ def on_select_segm_dir():
     pb_embed_sel = Container(widgets=[on_select_embed], layout="horizontal", labels=False)
     pb_seg_segm = Container(widgets=[on_select_segm, on_select_segm_dir], layout="horizontal", labels=False)
 
-    re_halo_x = SpinBox(value=0, max=10000, label="Halo x")
-    re_halo_y = SpinBox(value=0, max=10000, label="Halo y")
-    re_tile_x = SpinBox(value=0, max=10000, label="Tile x")
-    re_tile_y = SpinBox(value=0, max=10000, label="Tile y")
-    cb_model = ComboBox(value="vit_h", choices=["vit_h", "vit_l", "vit_b"], label="Model Type")
-
     @magicgui.magicgui(call_button="3d annotator", labels=False)
     def on_start():
         try:
@@ -240,6 +256,7 @@ def on_start():
     sub_widget = Container(widgets=[Container(widgets=[on_start], layout="horizontal", labels=False),
                                     pb_img_sel, lbl_opt, pb_embed_sel, pb_seg_segm, le_file_key_img, le_file_key_segm,
                                     cb_model, re_tile_x, re_tile_y, re_halo_x, re_halo_y])
+    sub_widget.root_native_widget.setWindowTitle("Segment Anything for Microscopy")
     main_widget.close()
     sub_widget.show()
 
@@ -250,6 +267,12 @@ def _on_series():
     config_dict["args"] = {}
     args = config_dict["args"]
 
+    re_halo_x = SpinBox(value=0, max=10000, label="Halo x")
+    re_halo_y = SpinBox(value=0, max=10000, label="Halo y")
+    re_tile_x = SpinBox(value=0, max=10000, label="Tile x")
+    re_tile_y = SpinBox(value=0, max=10000, label="Tile y")
+    cb_model = ComboBox(value="vit_h_lm", choices=get_model_names(), label="Model Type")
+
     @magicgui.magicgui(call_button="Select input directory", labels=False)
     def on_select_input_dir():
         try:
@@ -288,19 +311,14 @@ def on_select_embed():
                 show_error("Precompute embeddings file does not exist or has wrong file extension.")
                 return
             args["embedding_path"] = path
+            _set_embeddings_file_attributes(path, cb_model, re_tile_x, re_tile_y)
         except Exception as e:
             show_error(str(e))
 
     pb_input_sel = Container(widgets=[on_select_input_dir], layout="horizontal", labels=False)
     pb_output_sel = Container(widgets=[on_select_output_dir], layout="horizontal", labels=False)
     pb_embed_sel = Container(widgets=[on_select_embed], layout="horizontal", labels=False)
 
-    re_halo_x = SpinBox(value=0, max=10000, label="Halo x")
-    re_halo_y = SpinBox(value=0, max=10000, label="Halo y")
-    re_tile_x = SpinBox(value=0, max=10000, label="Tile x")
-    re_tile_y = SpinBox(value=0, max=10000, label="Tile y")
-    cb_model = ComboBox(value="vit_h", choices=["vit_h", "vit_l", "vit_b"], label="Model Type")
-
     @magicgui.magicgui(call_button="Image series annotator", labels=False)
     def on_start():
         try:
@@ -329,6 +347,7 @@ def on_start():
     sub_widget = Container(widgets=[Container(widgets=[on_start], layout="horizontal", labels=False),
                                     pb_input_sel, pb_output_sel, lbl_opt, pb_embed_sel, cb_model, re_tile_x,
                                     re_tile_y, re_halo_x, re_halo_y])
+    sub_widget.root_native_widget.setWindowTitle("Segment Anything for Microscopy")
     main_widget.close()
     sub_widget.show()
 
@@ -341,6 +360,12 @@ def _on_tracking():
     le_file_key_img = LineEdit(value="*", label="File key input")
     le_file_key_segm = LineEdit(value="*", label="File key segmentation")
 
+    re_halo_x = SpinBox(value=0, max=10000, label="Halo x")
+    re_halo_y = SpinBox(value=0, max=10000, label="Halo y")
+    re_tile_x = SpinBox(value=0, max=10000, label="Tile x")
+    re_tile_y = SpinBox(value=0, max=10000, label="Tile y")
+    cb_model = ComboBox(value="vit_h_lm", choices=get_model_names(), label="Model Type")
+
     @magicgui.magicgui(call_button="Select images", labels=False)
     def on_select_image():
         try:
@@ -379,6 +404,7 @@ def on_select_embed():
                 show_error("Precompute embeddings file does not exist or has wrong file extension.")
                 return
             args["embedding_path"] = path
+            _set_embeddings_file_attributes(path, cb_model, re_tile_x, re_tile_y)
         except Exception as e:
             show_error(str(e))
 
@@ -412,12 +438,6 @@ def on_select_result_dir():
     pb_embed_sel = Container(widgets=[on_select_embed], layout="horizontal", labels=False)
     pb_seg_sel = Container(widgets=[on_select_results, on_select_result_dir], layout="horizontal", labels=False)
 
-    re_halo_x = SpinBox(value=0, max=10000, label="Halo x")
-    re_halo_y = SpinBox(value=0, max=10000, label="Halo y")
-    re_tile_x = SpinBox(value=0, max=10000, label="Tile x")
-    re_tile_y = SpinBox(value=0, max=10000, label="Tile y")
-    cb_model = ComboBox(value="vit_h", choices=["vit_h", "vit_l", "vit_b"], label="Model Type")
-
     @magicgui.magicgui(call_button="Tracking annotator", labels=False)
     def on_start():
         try:
@@ -444,6 +464,7 @@ def on_start():
     sub_widget = Container(widgets=[Container(widgets=[on_start], layout="horizontal", labels=False),
                                     pb_img_sel, lbl_opt, pb_embed_sel, pb_seg_sel, le_file_key_img, le_file_key_segm,
                                     cb_model, re_tile_x, re_tile_y, re_halo_x, re_halo_y])
+    sub_widget.root_native_widget.setWindowTitle("Segment Anything for Microscopy")
     main_widget.close()
     sub_widget.show()
 
@@ -460,6 +481,7 @@ def annotator():
     sub_container2 = Container(widgets=[_on_3d, _on_tracking], labels=False)
     sub_container3 = Container(widgets=[sub_container1, sub_container2], layout="horizontal", labels=False)
     main_widget = Container(widgets=[Label(value="Segment Anything for Microscopy"), sub_container3], labels=False)
+    main_widget.root_native_widget.setWindowTitle("Segment Anything for Microscopy")
     main_widget.show(run=True)
 
     if config_dict["workflow"] == "2d":

micro_sam/sam_annotator/gui_utils.py

@@ -1,8 +1,12 @@
 import os
+import magicgui
 from shutil import rmtree
 from typing import Union
 
-from PyQt5 import QtCore, QtWidgets
+from magicgui.widgets import Container
+from magicgui.application import use_app
+
+from PyQt5 import QtWidgets
 
 
 def show_wrong_file_warning(file_path: Union[str, os.PathLike]) -> Union[str, os.PathLike]:
@@ -20,37 +24,62 @@ def show_wrong_file_warning(file_path: Union[str, os.PathLike]) -> Union[str, os
     Returns:
         Path to a file (new or old) depending on user decision
     """
-    msgbox = QtWidgets.QMessageBox()
-    msgbox.setWindowFlags(QtCore.Qt.CustomizeWindowHint | QtCore.Qt.WindowTitleHint)
-    msgbox.setWindowTitle("Warning")
-    msgbox.setText("The input data does not match the embeddings file.")
-    ignore_btn = msgbox.addButton("Ignore", QtWidgets.QMessageBox.RejectRole)
-    overwrite_btn = msgbox.addButton("Overwrite file", QtWidgets.QMessageBox.DestructiveRole)
-    select_btn = msgbox.addButton("Select different file", QtWidgets.QMessageBox.AcceptRole)
-    create_btn = msgbox.addButton("Create new file", QtWidgets.QMessageBox.AcceptRole)
-    msgbox.setDefaultButton(create_btn)
-
-    msgbox.exec()
-    msgbox.clickedButton()
-    if msgbox.clickedButton() == ignore_btn:
-        return file_path
-    elif msgbox.clickedButton() == overwrite_btn:
+#    q_app = QtWidgets.QApplication([])
+#    msgbox = QtWidgets.QMessageBox()
+#    msgbox.setWindowFlags(QtCore.Qt.CustomizeWindowHint | QtCore.Qt.WindowTitleHint)
+#    msgbox.setWindowTitle("Warning")
+#    msgbox.setText("The input data does not match the embeddings file.")
+#    create_btn = msgbox.addButton("Create new file", QtWidgets.QMessageBox.AcceptRole)
+
+    msg_box = None
+    new_path = {"value": ""}
+
+    @magicgui.magicgui(call_button="Ignore", labels=False)
+    def _ignore():
+        msg_box.close()
+        new_path["value"] = file_path
+
+    @magicgui.magicgui(call_button="Overwrite file", labels=False)
+    def _overwrite():
+        msg_box.close()
         rmtree(file_path)
-        return file_path
-    elif msgbox.clickedButton() == create_btn:
+        new_path["value"] = file_path
+
+    @magicgui.magicgui(call_button="Create new file", labels=False)
+    def _create():
+        msg_box.close()
         # unfortunately there exists no dialog to create a directory so we have
         # to use "create new file" dialog with some adjustments.
         dialog = QtWidgets.QFileDialog(None)
         dialog.setFileMode(QtWidgets.QFileDialog.AnyFile)
         dialog.setOption(QtWidgets.QFileDialog.ShowDirsOnly)
         dialog.setNameFilter("Archives (*.zarr)")
-        new_path = ""
-        while os.path.splitext(new_path)[1] != ".zarr":
-            dialog.exec()
-            new_path = dialog.selectedFiles()[0]
-        os.makedirs(new_path)
-        return(new_path)
-    elif msgbox.clickedButton() == select_btn:
-        return QtWidgets.QFileDialog.getExistingDirectory(
-            None, "Open a folder", os.path.split(file_path)[0], QtWidgets.QFileDialog.ShowDirsOnly
-        )
+        try_cnt = 0
+        while os.path.splitext(new_path["value"])[1] != ".zarr":
+            if try_cnt > 3:
+                new_path["value"] = file_path
+                return
+            dialog.exec_()
+            res = dialog.selectedFiles()
+            new_path["value"] = res[0] if len(res) > 0 else ""
+            try_cnt += 1
+        os.makedirs(new_path["value"])
+
+    @magicgui.magicgui(call_button="Select different file", labels=False)
+    def _select():
+        msg_box.close()
+        try_cnt = 0
+        while not os.path.exists(new_path["value"]):
+            if try_cnt > 3:
+                new_path["value"] = file_path
+                return
+            new_path["value"] = QtWidgets.QFileDialog.getExistingDirectory(
+                None, "Open a folder", os.path.split(file_path)[0], QtWidgets.QFileDialog.ShowDirsOnly
+            )
+            try_cnt += 1
+
+    msg_box = Container(widgets=[_select, _ignore, _overwrite, _create], layout='horizontal', labels=False)
+    msg_box.root_native_widget.setWindowTitle("The input data does not match the embeddings file")
+    msg_box.show(run=True)
+    use_app().quit()
+    return new_path["value"]

micro_sam/util.py

@@ -6,7 +6,7 @@
 import os
 import warnings
 from shutil import copyfileobj
-from typing import Any, Callable, Dict, Optional, Tuple
+from typing import Any, Callable, Dict, Optional, Tuple, Iterable
 
 import imageio.v3 as imageio
 import numpy as np
@@ -135,11 +135,16 @@ def get_sam_model(
     sam = sam_model_registry[model_type_](checkpoint=checkpoint)
     sam.to(device=device)
     predictor = SamPredictor(sam)
+    predictor.model_type = model_type
     if return_sam:
         return predictor, sam
     return predictor
 
 
+def get_model_names() -> Iterable:
+    return _MODEL_URLS.keys()
+
+
 def _to_image(input_):
     # we require the input to be uint8
     if input_.dtype != np.dtype("uint8"):
@@ -392,19 +397,22 @@ def precompute_image_embeddings(
         data_signature = _compute_data_signature(input_)
 
         f = zarr.open(save_path, "a")
-        if "input_size" in f.attrs:  # we have computed the embeddings already
-
-            # data signature does not match or is not in the file
-            if "data_signature" not in f.attrs or f.attrs["data_signature"] != data_signature:
-                warnings.warn("Embeddings file is invalid. Please recompute embeddings in a new file.")
-                if wrong_file_callback is not None:
-                    save_path = wrong_file_callback(save_path)
-                f = zarr.open(save_path, "a")
-                if "data_signature" not in f.attrs:
-                    f.attrs["data_signature"] = data_signature
-
-        else:  # embeddings have not yet been computed
-            f.attrs["data_signature"] = data_signature
+        key_vals = [("data_signature", data_signature),
+                    ("tile_shape", tile_shape), ("model_type", predictor.model_type)]
+        for key, val in key_vals:
+            if "input_size" in f.attrs:  # we have computed the embeddings already
+                # key signature does not match or is not in the file
+                if key not in f.attrs or f.attrs[key] != val:
+                    warnings.warn(f"Embeddings file is invalid due to unmatching {key}. \
+                        Please recompute embeddings in a new file.")
+                    if wrong_file_callback is not None:
+                        save_path = wrong_file_callback(save_path)
+                        f = zarr.open(save_path, "a")
+                    break
+
+        for key, val in key_vals:
+            if key not in f.attrs:
+                f.attrs[key] = val
 
     if ndim == 2:
         image_embeddings = _compute_2d(input_, predictor) if save_path is None else\