Skip to content

Commit 55a0f41

Browse files
committed
issue 40 add input data signature to embeddings file.
add dialogs for warning message in case of signature missmatch. add dialogs for selecting an existing file or for creating a new one.
1 parent fdf6a2b commit 55a0f41

File tree

1 file changed

+62
-1
lines changed

1 file changed

+62
-1
lines changed

micro_sam/util.py

Lines changed: 62 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,15 @@
11
import hashlib
22
import os
3-
from shutil import copyfileobj
3+
import warnings
4+
from shutil import copyfileobj, rmtree
45

56
import numpy as np
67
import requests
78
import torch
89
import vigra
910
import zarr
1011

12+
from PyQt5 import QtCore, QtWidgets
1113
from elf.io import open_file
1214
from nifty.tools import blocking
1315
from skimage.measure import regionprops
@@ -314,6 +316,51 @@ def _precompute_3d(input_, predictor, save_path, lazy_loading, tile_shape=None,
314316
}
315317
return image_embeddings
316318

319+
def show_wrong_file_warning(file_path):
320+
"""If the data signature does not match to the signature, user will can choose from the following options in this dialog:
321+
- Ignore: continue with input file (return file_path).
322+
- Overwrite: delete file_path and recompute the embeddings at same location.
323+
- Select a different file
324+
- Select a new file
325+
326+
Arguments:
327+
file_path (string or os.path): path of the problematic file
328+
329+
Returns:
330+
string or os.path: path to a file (new or old) depending on user decision
331+
"""
332+
msgbox = QtWidgets.QMessageBox()
333+
msgbox.setWindowFlags(QtCore.Qt.CustomizeWindowHint | QtCore.Qt.WindowTitleHint)
334+
msgbox.setWindowTitle("Warning")
335+
msgbox.setText('The input data does not match the embeddings file.')
336+
ignore_btn = msgbox.addButton("Ignore" , QtWidgets.QMessageBox.RejectRole)
337+
overwrite_btn = msgbox.addButton("Overwrite file" , QtWidgets.QMessageBox.DestructiveRole)
338+
select_btn = msgbox.addButton("Select different file" ,QtWidgets.QMessageBox.AcceptRole)
339+
create_btn = msgbox.addButton("Create new file" ,QtWidgets.QMessageBox.AcceptRole)
340+
msgbox.setDefaultButton(create_btn)
341+
342+
msgbox.exec()
343+
msgbox.clickedButton()
344+
if msgbox.clickedButton() == ignore_btn:
345+
return file_path
346+
elif msgbox.clickedButton() == overwrite_btn:
347+
rmtree(file_path)
348+
return file_path
349+
elif msgbox.clickedButton() == create_btn:
350+
# unfortunately there exists no dialog to create a directory so we have to use "create new file" dialog with some adjustments.
351+
dialog = QtWidgets.QFileDialog(None)
352+
dialog.setFileMode(QtWidgets.QFileDialog.AnyFile)
353+
dialog.setOption(QtWidgets.QFileDialog.ShowDirsOnly)
354+
dialog.setNameFilter("Archives (*.zarr)")
355+
new_path = ""
356+
while os.path.splitext(new_path)[1] != ".zarr":
357+
dialog.exec()
358+
new_path = dialog.selectedFiles()[0]
359+
os.makedirs(new_path)
360+
return(new_path)
361+
elif msgbox.clickedButton() == select_btn:
362+
return QtWidgets.QFileDialog.getExistingDirectory(None, "Open a folder", os.path.split(file_path)[0], QtWidgets.QFileDialog.ShowDirsOnly)
363+
317364

318365
def precompute_image_embeddings(
319366
predictor, input_, save_path=None, lazy_loading=False, ndim=None, tile_shape=None, halo=None
@@ -337,6 +384,20 @@ def precompute_image_embeddings(
337384
ndim = input_.ndim if ndim is None else ndim
338385
if tile_shape is not None:
339386
assert save_path is not None, "Tiled prediction is only supported when the embeddings are saved to file."
387+
388+
if save_path is not None:
389+
data_signature = hashlib.sha1(input_.tobytes()).hexdigest()
390+
391+
f = zarr.open(save_path, "a")
392+
if "input_size" in f.attrs:
393+
if "data_signature" not in f.attrs or f.attrs["data_signature"] != data_signature:
394+
warnings.warn("Embeddings file is invalid. Please recompute embeddings to new file.")
395+
save_path = show_wrong_file_warning(save_path)
396+
f = zarr.open(save_path, "a")
397+
if "data_signature" not in f.attrs:
398+
f.attrs["data_signature"] = data_signature
399+
else:
400+
f.attrs["data_signature"] = data_signature
340401

341402
if ndim == 2:
342403
image_embeddings = _compute_2d(input_, predictor) if save_path is None else\

0 commit comments

Comments
 (0)