Merge branch 'dino' of https://github.com/mouseland/cellpose into dino

Author: carsen-stringer

cellpose/cli.py

@@ -65,9 +65,9 @@ def get_arg_parser():
 
     # model settings
     model_args = parser.add_argument_group("Model Arguments")
-    model_args.add_argument("--pretrained_model", required=False, default="cpsam",
+    model_args.add_argument("--pretrained_model", required=False, default="cpdino",
                             type=str,
-                            help="model to use for running or starting training")
+                            help="path to model for segmentation or starting training, or builtin model name: cpdino, cpsam_v2, cpdino-vitb, or cpsam")
     model_args.add_argument(
         "--add_model", required=False, default=None, type=str,
         help="model path to copy model to hidden .cellpose folder for using in GUI/CLI")

cellpose/gui/gui.py

@@ -110,6 +110,8 @@ def __init__(self, image=None, logger=None):
         # MainW init
         MainW.__init__(self, image=image, logger=logger)
 
+        self.win.scene().sigMouseClicked.connect(self.plot_clicked)
+
         # add gradZ view
         self.ViewDropDown.insertItem(3, "gradZ")
 
@@ -289,9 +291,10 @@ def add_mask(self, points=None, color=(100, 200, 50), dense=True):
                     ar, ac = np.nonzero(mask)
                     ar, ac = ar + vr.min() - 2, ac + vc.min() - 2
                     # get dense outline
-                    contours = cv2.findContours(mask, cv2.RETR_EXTERNAL,
-                                                cv2.CHAIN_APPROX_NONE)
-                    pvc, pvr = contours[-2][0].squeeze().T
+                    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL,
+                                                   cv2.CHAIN_APPROX_NONE)
+                    contour = contours[np.argmax([c.size for c in contours])]
+                    pvc, pvr = contour.squeeze().T
                     vr, vc = pvr + vr.min() - 2, pvc + vc.min() - 2
                     # concatenate all points
                     ar, ac = np.hstack((np.vstack((vr, vc)), np.vstack((ar, ac))))
@@ -305,9 +308,9 @@ def add_mask(self, points=None, color=(100, 200, 50), dense=True):
                         # compute outline of new mask
                         mask = np.zeros((np.ptp(ar) + 4, np.ptp(ac) + 4), "uint8")
                         mask[ar - ar.min() + 2, ac - ac.min() + 2] = 1
-                        contours = cv2.findContours(mask, cv2.RETR_EXTERNAL,
-                                                    cv2.CHAIN_APPROX_NONE)
-                        pvc, pvr = contours[-2][0].squeeze().T
+                        contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
+                        contour = contours[np.argmax([c.size for c in contours])]
+                        pvc, pvr = contour.squeeze().T
                         vr, vc = pvr + ar.min() - 2, pvc + ac.min() - 2
                     ars = np.concatenate((ars, ar), axis=0)
                     acs = np.concatenate((acs, ac), axis=0)
@@ -560,13 +563,7 @@ def plot_clicked(self, event):
         if event.button()==QtCore.Qt.LeftButton \
                 and not event.modifiers() & (QtCore.Qt.ShiftModifier | QtCore.Qt.AltModifier)\
                 and not self.removing_region:
-            if event.double():
-                try:
-                    self.p0.setYRange(0, self.Ly + self.pr)
-                except:
-                    self.p0.setYRange(0, self.Ly)
-                self.p0.setXRange(0, self.Lx)
-            elif self.loaded and not self.in_stroke:
+            if self.loaded and not self.in_stroke:
                 if self.orthobtn.isChecked():
                     items = self.win.scene().items(event.scenePos())
                     for x in items:
@@ -621,6 +618,11 @@ def keyPressEvent(self, event):
                         self.scroll.setValue(self.currentZ)
                         event.accept()
                         return
+
+            # when in stroke, allow escaping out of drawing
+            else: 
+                if event.key() == QtCore.Qt.Key_Escape:
+                    self.layer.end_stroke(keep_stroke=False)
         if event.key() == QtCore.Qt.Key_Minus or event.key() == QtCore.Qt.Key_Equal:
             self.p0.keyPressEvent(event)
             event.accept()

cellpose/gui/gui3d.py

@@ -10,6 +10,7 @@
 import numpy as np
 import pathlib, os
 
+from cellpose import models
 from cellpose.gui.io import _save_sets
 
 
@@ -111,32 +112,6 @@ def setup(self):
         )
 
 
-# def create_channel_choose():
-#     # choose channel
-#     ChannelChoose = [QComboBox(), QComboBox()]
-#     ChannelLabels = []
-#     ChannelChoose[0].addItems(["gray", "red", "green", "blue"])
-#     ChannelChoose[1].addItems(["none", "red", "green", "blue"])
-#     cstr = ["chan to segment:", "chan2 (optional): "]
-#     for i in range(2):
-#         ChannelLabels.append(QLabel(cstr[i]))
-#         if i == 0:
-#             ChannelLabels[i].setToolTip(
-#                 "this is the channel in which the cytoplasm or nuclei exist \
-#             that you want to segment")
-#             ChannelChoose[i].setToolTip(
-#                 "this is the channel in which the cytoplasm or nuclei exist \
-#             that you want to segment")
-#         else:
-#             ChannelLabels[i].setToolTip(
-#                 "if <em>cytoplasm</em> model is chosen, and you also have a \
-#             nuclear channel, then choose the nuclear channel for this option")
-#             ChannelChoose[i].setToolTip(
-#                 "if <em>cytoplasm</em> model is chosen, and you also have a \
-#             nuclear channel, then choose the nuclear channel for this option")
-
-#     return ChannelChoose, ChannelLabels
-
 def unsilence_exceptions(func):
     """ Wrapper to unsilence Qt exceptions and re-raise them """
     def wrapper(*args, **kwargs):
@@ -148,19 +123,6 @@ def wrapper(*args, **kwargs):
             logger.debug(''.join(traceback.format_exception(type(e), e, e.__traceback__)))
     return wrapper
 
-class ModelButton(QPushButton):
-
-    def __init__(self, parent, model_name, text):
-        super().__init__()
-        self.setEnabled(False)
-        self.setText(text)
-        self.setFont(parent.boldfont)
-        self.clicked.connect(lambda: self.press(parent))
-        self.model_name = "cpsam"
-
-    def press(self, parent):
-        parent.compute_segmentation(model_name="cpsam")
-
 
 class FilterButton(QPushButton):
 
@@ -428,7 +390,7 @@ def niter(self):
 
 class TrainWindow(QDialog):
 
-    def __init__(self, parent, model_strings):
+    def __init__(self, parent):
         super().__init__(parent)
         self.setGeometry(100, 100, 900, 550)
         self.setWindowTitle("train settings")
@@ -446,7 +408,7 @@ def __init__(self, parent, model_strings):
         # choose initial model
         yoff += 1
         self.ModelChoose = QComboBox()
-        self.ModelChoose.addItems(model_strings)
+        self.ModelChoose.addItems(models.MODEL_NAMES)
         self.ModelChoose.setFixedWidth(150)
         self.ModelChoose.setCurrentIndex(parent.training_params["model_index"])
         self.l0.addWidget(self.ModelChoose, yoff, 1, 1, 1)
@@ -667,6 +629,7 @@ def mouseClickEvent(self, ev):
             if y >= 0 and y < self.parent.Ly and x >= 0 and x < self.parent.Lx:
                 if ev.button() == QtCore.Qt.LeftButton and not ev.double():
                     idx = self.parent.cellpix[self.parent.currentZ][y, x]
+                    self.parent.logger.debug(f'clicked on idx: {idx}')
                     if idx > 0:
                         if ev.modifiers() & QtCore.Qt.ControlModifier:
                             # delete mask selected
@@ -686,9 +649,17 @@ def mouseClickEvent(self, ev):
                     elif self.parent.masksOn and not self.parent.deleting_multiple:
                         self.parent.unselect_cell()
 
+    @unsilence_exceptions
+    def mouseDoubleClickEvent(self, ev) -> None:
+        ev.accept()
+        return
+
     @unsilence_exceptions
     def mouseDragEvent(self, ev):
-        ev.ignore()
+        if ev.button() == QtCore.Qt.RightButton:
+            ev.accept()
+        else:
+            ev.ignore()
         return
 
     @unsilence_exceptions
@@ -717,6 +688,9 @@ def is_at_start(self, pos):
         # first check if you ever left the start
         if len(self.parent.current_stroke) > 3:
             stroke = np.array(self.parent.current_stroke)
+            stroke = stroke[stroke[:, 0] == self.parent.currentZ]
+            if len(stroke) <= 3:
+                return False
             dist = (((stroke[1:, 1:] -
                       stroke[:1, 1:][np.newaxis, :, :])**2).sum(axis=-1))**0.5
             dist = dist.flatten()
@@ -731,10 +705,19 @@ def is_at_start(self, pos):
             else:
                 return False
 
-    def end_stroke(self):
+    def end_stroke(self, keep_stroke=True):
         if hasattr(self, 'scatter') and self.scatter is not None:
             if self.scatter.scene() == self.parent.layer.scene():
                 self.parent.p0.removeItem(self.scatter)
+        if not keep_stroke:
+            if not self.parent.stroke_appended:
+                self.parent.strokes.append(self.parent.current_stroke)
+                self.parent.stroke_appended = True
+            self.parent.remove_stroke(delete_points=False)
+            self.parent.current_stroke = [s for s in self.parent.current_stroke
+                                          if s[0] != self.parent.currentZ]
+            self.parent.in_stroke = False
+            return
         if not self.parent.stroke_appended:
             self.parent.strokes.append(self.parent.current_stroke)
             self.parent.stroke_appended = True

cellpose/gui/guiparts.py

@@ -41,23 +41,22 @@ def _add_model(parent, filename=None, load_model=True):
 
     for ind, model_string in enumerate(parent.model_strings[:-1]):
         if model_string == fname:
-            _remove_model(parent, ind=ind + 1, verbose=False)
+            _remove_model(parent, ind=ind + 1)
 
     parent.ModelChooseC.setCurrentIndex(len(parent.model_strings))
     if load_model:
         parent.model_choose(custom=True)
 
 
-def _remove_model(parent, ind=None, verbose=True):
+def _remove_model(parent, ind=None):
     if ind is None:
         ind = parent.ModelChooseC.currentIndex()
     if ind > 0:
-        ind -= 1
-        parent.ModelChooseC.removeItem(ind + 1)
-        del parent.model_strings[ind]
-        # remove model from txt path
         modelstr = parent.ModelChooseC.currentText()
+        parent.ModelChooseC.removeItem(ind)
+        # remove model from txt path
         remove_model(modelstr)
+        parent.model_strings.remove(modelstr)
         if len(parent.model_strings) > 0:
             parent.ModelChooseC.setCurrentIndex(len(parent.model_strings))
         else:
@@ -329,7 +328,11 @@ def _load_seg(parent, filename=None, image=None, image_file=None, load_3D=False)
             parent.ismanual = dat["ismanual"]
 
     if "current_channel" in dat:
-        parent.color = (dat["current_channel"] + 2) % 5
+        # saved as int, loading int is built into .color property
+        color = dat['current_channel']
+        if isinstance(color, tuple):
+            color = color[0]
+        parent.color = color
 
     if "flows" in dat:
         parent.flows = dat["flows"]
@@ -355,6 +358,7 @@ def _load_seg(parent, filename=None, image=None, image_file=None, load_3D=False)
 
     parent.enable_buttons()
     parent.update_layer()
+    parent.update_plot()
     del dat
     gc.collect()
 
@@ -551,6 +555,8 @@ def _save_sets(parent):
     flow_threshold = parent.segmentation_settings.flow_threshold
     cellprob_threshold = parent.segmentation_settings.cellprob_threshold
 
+    # use ints instead of strings for backwards compatibility with old _seg.npy files
+    color_int = list(parent.RGBDropDown.name_map.keys()).index(parent.color)
     if parent.NZ > 1:
         dat = {
             "outlines":
@@ -559,7 +565,7 @@ def _save_sets(parent):
                 parent.cellcolors[1:],
             "masks":
                 parent.cellpix,
-            "current_channel": (parent.color - 2) % 5,
+            "current_channel": color_int,
             "filename":
                 parent.filename,
             "flows":

cellpose/gui/io.py

@@ -79,7 +79,7 @@ def logger_setup(cp_path=".cellpose", logfile_name="run.log", stdout_file_replac
     log_file = cp_dir.joinpath(logfile_name)
     try:
         log_file.unlink()
-    except:
+    except FileNotFoundError:
         print('creating new log file')
     logfile_fh = logging.FileHandler(log_file)
     if stdout_file_replacement is not None:
@@ -348,6 +348,9 @@ def remove_model(filename, delete=False):
     filename = os.path.split(filename)[-1]
     from . import models
     model_strings = models.get_user_models()
+    if filename not in model_strings:
+        raise ValueError(f'filename not found: {filename}')
+    model_strings.remove(filename)
     if len(model_strings) > 0:
         with open(models.MODEL_LIST_PATH, "w") as textfile:
             for fname in model_strings:
@@ -356,10 +359,10 @@ def remove_model(filename, delete=False):
         # write empty file
         textfile = open(models.MODEL_LIST_PATH, "w")
         textfile.close()
-    print(f"{filename} removed from custom model list")
+    io_logger.info(f"{filename} removed from custom model list")
     if delete:
-        os.remove(os.fspath(models.MODEL_DIR.joinpath(fname)))
-        print("model deleted")
+        os.remove(os.fspath(models.MODEL_DIR.joinpath(filename)))
+        io_logger.info(f"{filename} model deleted from disk")
 
 
 def add_model(filename):

cellpose/io.py

@@ -24,7 +24,7 @@
 _MODEL_DIR_DEFAULT = Path.home().joinpath(".cellpose", "models")
 MODEL_DIR = Path(_MODEL_DIR_ENV) if _MODEL_DIR_ENV else _MODEL_DIR_DEFAULT
 
-MODEL_NAMES = ["cpsam", "cpsam_v2", "cpdino", "cpdino-vitb"]
+MODEL_NAMES = ["cpdino", "cpsam_v2", "cpdino-vitb", "cpsam"]
 
 MODEL_LIST_PATH = os.fspath(MODEL_DIR.joinpath("gui_models.txt"))
 

cellpose/models.py

@@ -873,12 +873,13 @@ def random_rotate_and_resize(imgs, lbls=None, scale_range=1., bsize=256, rescale
         if rescale is not None:
             scale *= 1. / rescale[n]
         scale_per_image[n] = scale
-        angle = 360 * np.random.rand() - 180
+        angle = np.float32(360 * np.random.rand() - 180)
         Ly, Lx = imgs[n].shape[-2:]
-        dyx = np.maximum(0, np.array([Ly * scale - yx[0], Lx * scale - yx[1]]))
+        dyx = np.maximum(0, np.array([Ly * scale - yx[0], Lx * scale - yx[1]])).astype("float32")
         flip = np.random.rand() > 0.5
-        translations = list(dyx * np.random.rand(2) * 1 - dyx * 0.5)
-        matrix = F._get_inverse_affine_matrix([0., 0.], angle, translations, scale, [0., 0.])
+        translations = list(dyx * np.random.rand(2).astype("float32") * 1 - dyx * 0.5)
+        matrix = F._get_inverse_affine_matrix(np.zeros(2, "float32"), angle, translations, 
+                                              scale, np.zeros(2, "float32"))
         theta = torch.tensor(matrix, device=device).reshape(1, 2, 3)
         grid = F_t._gen_affine_grid(theta, w=Lx, h=Ly, ow=yx[1], oh=yx[0])
 

cellpose/transforms.py

@@ -40,28 +40,28 @@ For multi-Z 3D data, please use the 3D version of the GUI:
 Using the GUI 
 ~~~~~~~~~~~~~~~~~~~~~~~
 
-The GUI serves two main functions:
+The GUI serves three main functions:
 
 1. Running the segmentation algorithm.
 2. Manually labelling data.
-3. (NEW) Fine-tuning a pretrained cellpose model on your own data.
+3. Fine-tuning a pretrained cellpose model on your own data.
 
 Main GUI mouse controls (works in all views):
 
 -  Pan = left-click + drag
 -  Zoom = scroll wheel (or +/= and - buttons)
--  Full view = double left-click
 -  Select mask = left-click on mask
 -  Delete mask = Ctrl (or Command on Mac) + left-click
 -  Merge masks = Alt + left-click (will merge last two)
 -  Start draw mask = right-click
--  End draw mask = right-click, or return to circle at beginning
+-  End draw mask = press esc, right-click, or return to circle at beginning
 
 Drawing masks 
 ~~~~~~~~~~~~~~~~~~~~~~
 
 Masks are started with right-click, then hover your mouse (do not hold it down), 
 and return it to the red circle to complete the mask. The mask should now be completed.
+At any point, you can escape drawing a mask with the Esc key.
 
 Overlaps in masks are NOT allowed. If you draw a mask on top of another
 mask, it is cropped so that it doesn't overlap with the old mask. Masks
@@ -146,6 +146,8 @@ Keyboard shortcuts
 +---------------------+-----------------------------------------------+            
 | =/+  // -           | zoom in // zoom out                           |
 +---------------------+-----------------------------------------------+
+| SPACE               | reset zoom                                    |
++---------------------+-----------------------------------------------+
 | CTRL+Z              | undo previously drawn mask/stroke             |
 +---------------------+-----------------------------------------------+
 | CTRL+0              | clear all masks                               |