Merge pull request #78 from computational-cell-analytics/refactor-gui

Refactor the wrong file dialog gui

Author: constantinpape

micro_sam/sam_annotator/annotator_2d.py

@@ -9,6 +9,7 @@
 from .. import util
 from .. import instance_segmentation
 from ..visualization import project_embeddings_for_visualization
+from .gui_utils import show_wrong_file_warning
 from .util import (
     clear_all_prompts, commit_segmentation_widget, create_prompt_menu,
     prompt_layer_to_boxes, prompt_layer_to_points, prompt_segmentation, toggle_label, LABEL_COLOR_CYCLE,
@@ -215,7 +216,8 @@ def annotator_2d(
     else:
         PREDICTOR = predictor
     IMAGE_EMBEDDINGS = util.precompute_image_embeddings(
-        PREDICTOR, raw, save_path=embedding_path, ndim=2, tile_shape=tile_shape, halo=halo
+        PREDICTOR, raw, save_path=embedding_path, ndim=2, tile_shape=tile_shape, halo=halo,
+        wrong_file_callback=show_wrong_file_warning
     )
 
     # we set the pre-computed image embeddings if we don't use tiling

micro_sam/sam_annotator/annotator_3d.py

@@ -8,6 +8,7 @@
 from .. import util
 from ..prompt_based_segmentation import segment_from_mask
 from ..visualization import project_embeddings_for_visualization
+from .gui_utils import show_wrong_file_warning
 from .util import (
     clear_all_prompts, commit_segmentation_widget, create_prompt_menu,
     prompt_layer_to_boxes, prompt_layer_to_points, prompt_segmentation,
@@ -195,7 +196,8 @@ def annotator_3d(
     global PREDICTOR, IMAGE_EMBEDDINGS
     PREDICTOR = util.get_sam_model(model_type=model_type)
     IMAGE_EMBEDDINGS = util.precompute_image_embeddings(
-        PREDICTOR, raw, save_path=embedding_path, tile_shape=tile_shape, halo=halo
+        PREDICTOR, raw, save_path=embedding_path, tile_shape=tile_shape, halo=halo,
+        wrong_file_callback=show_wrong_file_warning,
     )
 
     #

micro_sam/sam_annotator/annotator_tracking.py

@@ -12,6 +12,7 @@
 
 from .. import util
 from ..prompt_based_segmentation import segment_from_mask
+from .gui_utils import show_wrong_file_warning
 from .util import (
     create_prompt_menu, clear_all_prompts,
     prompt_layer_to_boxes, prompt_layer_to_points,
@@ -358,7 +359,8 @@ def annotator_tracking(
 
     PREDICTOR = util.get_sam_model(model_type=model_type)
     IMAGE_EMBEDDINGS = util.precompute_image_embeddings(
-        PREDICTOR, raw, save_path=embedding_path, tile_shape=tile_shape, halo=halo
+        PREDICTOR, raw, save_path=embedding_path, tile_shape=tile_shape, halo=halo,
+        wrong_file_callback=show_wrong_file_warning,
     )
 
     CURRENT_TRACK_ID = 1

micro_sam/sam_annotator/gui_utils.py

@@ -0,0 +1,54 @@
+import os
+from shutil import rmtree
+
+from PyQt5 import QtCore, QtWidgets
+
+
+def show_wrong_file_warning(file_path):
+    """If the data signature does not match to the signature,
+    the user can choose from the following options in this dialog:
+       - Ignore: continue with input file (return file_path).
+       - Overwrite: delete file_path and recompute the embeddings at same location.
+       - Select a different file
+       - Select a new file
+
+    Arguments:
+        file_path (string or os.path): path of the problematic file
+
+    Returns:
+        string or os.path: 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:
+        rmtree(file_path)
+        return file_path
+    elif msgbox.clickedButton() == create_btn:
+        # 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
+        )

micro_sam/util.py

@@ -1,15 +1,14 @@
 import hashlib
 import os
 import warnings
-from shutil import copyfileobj, rmtree
+from shutil import copyfileobj
 
 import numpy as np
 import requests
 import torch
 import vigra
 import zarr
 
-from PyQt5 import QtCore, QtWidgets
 from elf.io import open_file
 from nifty.tools import blocking
 from skimage.measure import regionprops
@@ -316,54 +315,12 @@ def _precompute_3d(input_, predictor, save_path, lazy_loading, tile_shape=None,
     }
     return image_embeddings
 
-def show_wrong_file_warning(file_path):
-    """If the data signature does not match to the signature, user will can choose from the following options in this dialog:
-       - Ignore: continue with input file (return file_path).
-       - Overwrite: delete file_path and recompute the embeddings at same location.
-       - Select a different file
-       - Select a new file
-       
-    Arguments:
-        file_path (string or os.path): path of the problematic file
-
-    Returns:
-        string or os.path: 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:
-        rmtree(file_path)
-        return file_path 
-    elif msgbox.clickedButton() == create_btn:
-        # 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)
-
 
 def precompute_image_embeddings(
-    predictor, input_, save_path=None, lazy_loading=False, ndim=None, tile_shape=None, halo=None
+    predictor, input_,
+    save_path=None, lazy_loading=False,
+    ndim=None, tile_shape=None, halo=None,
+    wrong_file_callback=None,
 ):
     """Compute the image embeddings (output of the encoder) for the input.
 
@@ -380,23 +337,31 @@ def precompute_image_embeddings(
         tile_shape [tuple] - shape of tiles for tiled prediction.
             By default prediction is run without tiling. (default: None)
         halo [tuple] - additional overlap of the tiles for tiled prediction. (default: None)
+        wrong_file_callback [callable] - function to call when an embedding file with wrong file signature
+            is passed. If none is given a wrong file signature will cause a warning.
+            If passed, the callback should have the signature 'def callback(save_path): return str',
+            where the return value is the (potentially updated) embedding save path (default: None)
     """
     ndim = input_.ndim if ndim is None else ndim
     if tile_shape is not None:
         assert save_path is not None, "Tiled prediction is only supported when the embeddings are saved to file."
-        
+
     if save_path is not None:
         data_signature = hashlib.sha1(input_.tobytes()).hexdigest()
-        
+
         f = zarr.open(save_path, "a")
-        if "input_size" in f.attrs:
+        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 to new file.")
-                save_path = show_wrong_file_warning(save_path)
+                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:
+                    f.attrs["data_signature"] = data_signature
+
+        else:  # embeddings have not yet been computed
             f.attrs["data_signature"] = data_signature
 
     if ndim == 2: