3D GUI: compute ortho cell outlines separately

cellpose/gui/gui.py

@@ -589,11 +589,12 @@ def make_buttons(self):
         w = 3
         for j in range(len(self.filter_text)):
             self.FilterButtons.append(
-                guiparts.FilterButton(self, self.filter_text[j]))
+                guiparts.FilterButton(self.medfont, self.filter_text[j]))
             self.filterBox_grid_layout.addWidget(self.FilterButtons[-1], widget_row, jj, 1, w)
             self.FilterButtons[-1].setFixedWidth(75)
             self.FilterButtons[-1].setToolTip(nett[j])
             self.FilterButtons[-1].setFont(self.medfont)
+            self.FilterButtons[-1].pressed_type.connect(self.filter_button_pressed)
             widget_row += 1 if j%2==1 else 0
             jj = 0 if j%2==1 else jj + w
 
@@ -777,8 +778,6 @@ def enable_buttons(self):
 
         for i in range(len(self.FilterButtons)):
             self.FilterButtons[i].setEnabled(True)
-        if self.load_3D:
-            self.FilterButtons[-2].setEnabled(False)
 
         self.newmodel.setEnabled(True)
         self.loadMasks.setEnabled(True)
@@ -1028,18 +1027,10 @@ def brush_choose(self):
     def clear_all(self):
         self.prev_selected = 0
         self.selected = 0
-        if self.restore and "upsample" in self.restore:
-            self.layerz = 0 * np.ones((self.Lyr, self.Lxr, 4), np.uint8)
-            self.cellpix = np.zeros((self.NZ, self.Lyr, self.Lxr), np.uint16)
-            self.outpix = np.zeros((self.NZ, self.Lyr, self.Lxr), np.uint16)
-            self.cellpix_resize = self.cellpix.copy()
-            self.outpix_resize = self.outpix.copy()
-            self.cellpix_orig = np.zeros((self.NZ, self.Ly0, self.Lx0), np.uint16)
-            self.outpix_orig = np.zeros((self.NZ, self.Ly0, self.Lx0), np.uint16)
-        else:
-            self.layerz = 0 * np.ones((self.Ly, self.Lx, 4), np.uint8)
-            self.cellpix = np.zeros((self.NZ, self.Ly, self.Lx), np.uint16)
-            self.outpix = np.zeros((self.NZ, self.Ly, self.Lx), np.uint16)
+        self.layerz = 0 * np.ones((self.Ly, self.Lx, 4), np.uint8)
+        self.cellpix = np.zeros((self.NZ, self.Ly, self.Lx), np.uint16)
+        self.outpix = np.zeros((self.NZ, self.Ly, self.Lx), np.uint16)
+        self.update_ortho_outpix()
 
         self.cellcolors = np.array([255, 255, 255])[np.newaxis, :]
         self.ncells.reset()
@@ -1126,6 +1117,8 @@ def remove_single_cell(self, idx):
         self.cellpix[self.cellpix > idx] -= 1
         self.outpix[self.outpix > idx] -= 1
 
+        self.update_ortho_outpix()
+
         if self.NZ == 1:
             self.removed_cell = [
                 self.ismanual[idx - 1], self.cellcolors[idx],
@@ -1494,48 +1487,27 @@ def draw_mask(self, z, ar, ac, vr, vc, color, idx=None):
         self.cellpix[z, vr, vc] = idx
         self.cellpix[z, ar, ac] = idx
         self.outpix[z, vr, vc] = idx
-        if self.restore and "upsample" in self.restore:
-            if self.resize:
-                self.cellpix_resize[z, vr, vc] = idx
-                self.cellpix_resize[z, ar, ac] = idx
-                self.outpix_resize[z, vr, vc] = idx
-                self.cellpix_orig[z, (vr / self.ratio).astype(int),
-                                  (vc / self.ratio).astype(int)] = idx
-                self.cellpix_orig[z, (ar / self.ratio).astype(int),
-                                  (ac / self.ratio).astype(int)] = idx
-                self.outpix_orig[z, (vr / self.ratio).astype(int),
-                                 (vc / self.ratio).astype(int)] = idx
-            else:
-                self.cellpix_orig[z, vr, vc] = idx
-                self.cellpix_orig[z, ar, ac] = idx
-                self.outpix_orig[z, vr, vc] = idx
-
-                # get upsampled mask
-                vrr = (vr.copy() * self.ratio).astype(int)
-                vcr = (vc.copy() * self.ratio).astype(int)
-                mask = np.zeros((np.ptp(vrr) + 4, np.ptp(vcr) + 4), np.uint8)
-                pts = np.stack((vcr - vcr.min() + 2, vrr - vrr.min() + 2),
-                               axis=-1)[:, np.newaxis, :]
-                mask = cv2.fillPoly(mask, [pts], (255, 0, 0))
-                arr, acr = np.nonzero(mask)
-                arr, acr = arr + vrr.min() - 2, acr + vcr.min() - 2
-                # get dense outline
-                contours = cv2.findContours(mask, cv2.RETR_EXTERNAL,
-                                            cv2.CHAIN_APPROX_NONE)
-                pvc, pvr = contours[-2][0].squeeze().T
-                vrr, vcr = pvr + vrr.min() - 2, pvc + vcr.min() - 2
-                # concatenate all points
-                arr, acr = np.hstack((np.vstack((vrr, vcr)), np.vstack((arr, acr))))
-                self.cellpix_resize[z, vrr, vcr] = idx
-                self.cellpix_resize[z, arr, acr] = idx
-                self.outpix_resize[z, vrr, vcr] = idx
+
+        self.update_ortho_outpix()
 
         if z == self.currentZ:
             self.layerz[ar, ac, :3] = color
             if self.masksOn:
                 self.layerz[ar, ac, -1] = self.opacity
             if self.outlinesOn:
                 self.layerz[vr, vc] = np.array(self.outcolor)
+
+    def update_ortho_outpix(self):
+        if self.NZ == 1: 
+            return
+
+        # calculate Y and X outlines: 
+        outlines_Y = masks_to_outlines(self.cellpix.transpose(1, 2, 0)).transpose(2, 0, 1)
+        self.outpix_Y = outlines_Y * self.cellpix
+
+        outlines_X = masks_to_outlines(self.cellpix.transpose(2, 1, 0)).transpose(2, 1, 0)
+        self.outpix_X = outlines_X * self.cellpix
+
 
     def compute_scale(self):
         # get diameter from gui
@@ -1570,15 +1542,6 @@ def draw_layer(self):
         else:
             self.Ly, self.Lx = self.Ly0, self.Lx0
 
-        if self.masksOn or self.outlinesOn:
-            if self.restore and "upsample" in self.restore:
-                if self.resize:
-                    self.cellpix = self.cellpix_resize.copy()
-                    self.outpix = self.outpix_resize.copy()
-                else:
-                    self.cellpix = self.cellpix_orig.copy()
-                    self.outpix = self.outpix_orig.copy()
-
         self.layerz = np.zeros((self.Ly, self.Lx, 4), np.uint8)
         if self.masksOn:
             self.layerz[..., :3] = self.cellcolors[self.cellpix[self.currentZ], :]
@@ -1671,7 +1634,7 @@ def compute_saturation_if_checked(self):
 
     def compute_saturation(self, return_img=False):
         norm = self.get_normalize_params()
-        print(norm)
+        self.logger.info(f'Normalization settings: {norm}')
         sharpen, smooth = norm["sharpen_radius"], norm["smooth_radius"]
         percentile = norm["percentile"]
         tile_norm = norm["tile_norm_blocksize"]
@@ -1716,6 +1679,8 @@ def compute_saturation(self, return_img=False):
             self.ViewDropDown.setCurrentIndex(self.ViewDropDown.count() - 1)
         else:
             img_norm = self.stack if self.restore is None or self.restore == "filter" else self.stack_filtered
+            self.ViewDropDown.model().item(self.ViewDropDown.count() - 1).setEnabled(False)
+            self.ViewDropDown.setCurrentIndex(0)
 
         if self.autobtn.isChecked():
             self.saturation = []
@@ -1747,7 +1712,7 @@ def compute_saturation(self, return_img=False):
                 else:
                     for n in range(self.NZ):
                         self.saturation[-1].append([0, 255.])
-            print(self.saturation[2][self.currentZ])
+            self.logger.debug(f'compute_saturation: saturation[2] levels at currentZ: {self.saturation[2][self.currentZ]}')
 
             if img_norm.shape[-1] == 1:
                 self.saturation.append(self.saturation[0])
@@ -1952,9 +1917,6 @@ def compute_segmentation(self, custom=False, model_name=None, load_model=True):
                     flows_new.append(np.zeros(flows[1][0].shape, dtype="uint8"))
 
             if not self.load_3D:
-                if self.restore and "upsample" in self.restore:
-                    self.Ly, self.Lx = self.Lyr, self.Lxr
-
                 if flows_new[0].shape[-3:-1] != (self.Ly, self.Lx):
                     self.flows = []
                     for j in range(len(flows_new)):
@@ -2005,3 +1967,28 @@ def compute_segmentation(self, custom=False, model_name=None, load_model=True):
                 self.recompute_masks = False
         except Exception as e:
             print("ERROR: %s" % e)
+
+    def filter_button_pressed(self, model_type: str):
+        if model_type == 'filter':
+            normalize_params = self.get_normalize_params()
+            if (normalize_params["sharpen_radius"] == 0 and
+                    normalize_params["smooth_radius"] == 0 and
+                    normalize_params["tile_norm_blocksize"] == 0):
+                print(
+                    "GUI_ERROR: no filtering settings on (use custom filter settings)")
+                self.restore = None
+                return
+            self.restore = "filter"
+
+        elif model_type == 'none':
+            self.set_normalize_params(
+                {
+                    'sharpen_radius' : 0,
+                    'smooth_radius' : 0,
+                    'tile_norm_blocksize' : 0,
+                }
+            )
+            self.restore = None
+        else: 
+            self.logger.error(f'Filter model type not known: {model_type}')
+        self.compute_saturation()

cellpose/gui/gui3d.py

@@ -376,6 +376,7 @@ def add_orthoviews(self):
         self.yortho = self.Ly // 2
         self.xortho = self.Lx // 2
         if self.NZ > 1:
+            self.update_ortho_outpix()
             self.update_ortho()
 
         self.win.addItem(self.pOrtho[0], 0, 1, rowspan=1, colspan=1)
@@ -429,8 +430,6 @@ def update_ortho(self):
                 self.pOrtho[0].setXRange(-self.dz / 3, self.dz * 2 + self.dz / 3)
                 self.pOrtho[1].setYRange(-self.dz / 3, self.dz * 2 + self.dz / 3)
             dztot = min(self.NZ, self.dz * 2)
-            y = self.yortho
-            x = self.xortho
             z = self.currentZ
             if dztot == self.NZ:
                 zmin, zmax = 0, self.NZ
@@ -445,6 +444,8 @@ def update_ortho(self):
                     zmin, zmax = z - self.dz, z + self.dz
             self.zc = z - zmin
             self.update_crosshairs()
+            y = self.yortho
+            x = self.xortho
             if self.view == 0 or self.view == 4:
                 for j in range(2):
                     if j == 0:
@@ -523,12 +524,13 @@ def update_ortho(self):
                         [255, 255, 255, self.opacity])
 
         if self.outlinesOn:
-            for j in range(2):
-                if j == 0:
-                    op = self.outpix[zmin:zmax, :, x].T
-                else:
-                    op = self.outpix[zmin:zmax, y]
-                self.layer_ortho[j][op > 0] = np.array(self.outcolor).astype("uint8")
+            # update j == 0 (YZ)
+            op = self.outpix_X[zmin:zmax, :, x].T
+            self.layer_ortho[0][op > 0] = np.array(self.outcolor).astype("uint8")
+
+            # update j == 1 (ZX)
+            op = self.outpix_Y[zmin:zmax, y, :]
+            self.layer_ortho[1][op > 0] = np.array(self.outcolor).astype("uint8")
 
         for j in range(2):
             self.layerOrtho[j].setImage(self.layer_ortho[j])

cellpose/gui/guiparts.py

@@ -150,32 +150,18 @@ def press(self, parent):
 
 class FilterButton(QPushButton):
 
-    def __init__(self, parent, text):
+    pressed_type = QtCore.Signal(str)
+
+    def __init__(self, font, text):
         super().__init__()
         self.setEnabled(False)
         self.model_type = text
         self.setText(text)
-        self.setFont(parent.medfont)
-        self.clicked.connect(lambda: self.press(parent))
+        self.setFont(font)
+        self.clicked.connect(self._emit_type)
 
-    def press(self, parent):
-        if self.model_type == "filter":
-            parent.restore = "filter"
-            normalize_params = parent.get_normalize_params()
-            if (normalize_params["sharpen_radius"] == 0 and
-                    normalize_params["smooth_radius"] == 0 and
-                    normalize_params["tile_norm_blocksize"] == 0):
-                print(
-                    "GUI_ERROR: no filtering settings on (use custom filter settings)")
-                parent.restore = None
-                return
-            parent.restore = self.model_type
-            parent.compute_saturation()
-        # elif self.model_type != "none":
-        #     parent.compute_denoise_model(model_type=self.model_type)
-        else:
-            parent.clear_restore()
-        # parent.set_restore_button()
+    def _emit_type(self):
+        self.pressed_type.emit(self.model_type)
 
 
 class ObservableVariable(QtCore.QObject):

cellpose/gui/io.py

@@ -10,6 +10,10 @@
 from ..models import normalize_default, MODEL_DIR, MODEL_LIST_PATH, get_user_models
 from ..utils import masks_to_outlines, outlines_list
 
+import logging 
+
+logger = logging.getLogger(__name__)
+
 try:
     import qtpy
     from qtpy.QtWidgets import QFileDialog
@@ -195,16 +199,13 @@ def _initialize_images(parent, image, load_3D=False):
     parent.layerz = 255 * np.ones((parent.Ly, parent.Lx, 4), "uint8")
     if hasattr(parent, "stack_filtered"):
         parent.Lyr, parent.Lxr = parent.stack_filtered.shape[-3:-1]
-    elif parent.restore and "upsample" in parent.restore:
-        parent.Lyr, parent.Lxr = int(parent.Ly * parent.ratio), int(parent.Lx *
-                                                                    parent.ratio)
     else:
         parent.Lyr, parent.Lxr = parent.Ly, parent.Lx
     parent.clear_all()
 
     if not hasattr(parent, "stack_filtered") and parent.restore:
         print("GUI_INFO: no 'img_restore' found, applying current settings")
-        parent.compute_restore()
+        parent.compute_saturation()
 
     if parent.autobtn.isChecked():
         if parent.restore is None or parent.restore != "filter":
@@ -286,7 +287,7 @@ def _load_seg(parent, filename=None, image=None, image_file=None, load_3D=False)
         parent.set_normalize_params(dat["normalize_params"])
 
     _initialize_images(parent, image, load_3D=load_3D)
-    print(parent.stack.shape)
+    logger.debug(f'Loaded segmentation image stake shape: {parent.stack.shape}')
 
     if "outlines" in dat:
         if isinstance(dat["outlines"], list):
@@ -415,58 +416,26 @@ def _masks_to_gui(parent, masks, outlines=None, colors=None):
         outlines = None
     masks = masks.astype(np.uint16) if masks.max() < 2**16 - 1 else masks.astype(
         np.uint32)
-    if parent.restore and "upsample" in parent.restore:
-        parent.cellpix_resize = masks.copy()
-        parent.cellpix = parent.cellpix_resize.copy()
-        parent.cellpix_orig = cv2.resize(
-            masks.squeeze(), (parent.Lx0, parent.Ly0),
-            interpolation=cv2.INTER_NEAREST)[np.newaxis, :, :]
-        parent.resize = True
-    else:
-        parent.cellpix = masks
+    parent.cellpix = masks
     if parent.cellpix.ndim == 2:
         parent.cellpix = parent.cellpix[np.newaxis, :, :]
-        if parent.restore and "upsample" in parent.restore:
-            if parent.cellpix_resize.ndim == 2:
-                parent.cellpix_resize = parent.cellpix_resize[np.newaxis, :, :]
-            if parent.cellpix_orig.ndim == 2:
-                parent.cellpix_orig = parent.cellpix_orig[np.newaxis, :, :]
 
     print(f"GUI_INFO: {masks.max()} masks found")
 
     # get outlines
     if outlines is None:  # parent.outlinesOn
         parent.outpix = np.zeros_like(parent.cellpix)
-        if parent.restore and "upsample" in parent.restore:
-            parent.outpix_orig = np.zeros_like(parent.cellpix_orig)
-        for z in range(parent.NZ):
-            outlines = masks_to_outlines(parent.cellpix[z])
-            parent.outpix[z] = outlines * parent.cellpix[z]
-            if parent.restore and "upsample" in parent.restore:
-                outlines = masks_to_outlines(parent.cellpix_orig[z])
-                parent.outpix_orig[z] = outlines * parent.cellpix_orig[z]
-            if z % 50 == 0 and parent.NZ > 1:
-                print("GUI_INFO: plane %d outlines processed" % z)
-        if parent.restore and "upsample" in parent.restore:
-            parent.outpix_resize = parent.outpix.copy()
+
+        outlines = masks_to_outlines(parent.cellpix)
+        parent.outpix = outlines * parent.cellpix
+
     else:
-        parent.outpix = outlines
-        if parent.restore and "upsample" in parent.restore:
-            parent.outpix_resize = parent.outpix.copy()
-            parent.outpix_orig = np.zeros_like(parent.cellpix_orig)
-            for z in range(parent.NZ):
-                outlines = masks_to_outlines(parent.cellpix_orig[z])
-                parent.outpix_orig[z] = outlines * parent.cellpix_orig[z]
-                if z % 50 == 0 and parent.NZ > 1:
-                    print("GUI_INFO: plane %d outlines processed" % z)
+        parent.outpix = outlines # set YX outlines
 
     if parent.outpix.ndim == 2:
         parent.outpix = parent.outpix[np.newaxis, :, :]
-        if parent.restore and "upsample" in parent.restore:
-            if parent.outpix_resize.ndim == 2:
-                parent.outpix_resize = parent.outpix_resize[np.newaxis, :, :]
-            if parent.outpix_orig.ndim == 2:
-                parent.outpix_orig = parent.outpix_orig[np.newaxis, :, :]
+
+    parent.update_ortho_outpix()
 
     parent.ncells.set(parent.cellpix.max())
     colors = parent.colormap[:parent.ncells.get(), :3] if colors is None else colors