Merge pull request #1461 from MouseLand/fix-sat

Fix sat

Author: carsen-stringer

cellpose/gui/gui.py

@@ -674,10 +674,10 @@ def level_change(self, r):
         if self.loaded:
             sval = self.sliders[r].value()
             self.saturation[r][self.currentZ] = sval
-            if not self.autobtn.isChecked():
-                for r in range(3):
-                    for i in range(len(self.saturation[r])):
-                        self.saturation[r][i] = self.saturation[r][self.currentZ]
+            # if not self.autobtn.isChecked():
+            for r in range(3):
+                for i in range(len(self.saturation[r])):
+                    self.saturation[r][i] = self.saturation[r][self.currentZ]
             self.update_plot()
 
     def keyPressEvent(self, event):
@@ -834,8 +834,7 @@ def enable_buttons(self):
     def disable_buttons_removeROIs(self):
         if len(self.model_strings) > 0:
             self.ModelButtonC.setEnabled(False)
-        for i in range(len(self.StyleButtons)):
-            self.StyleButtons[i].setEnabled(False)
+        self.ModelButtonB.setEnabled(False)
         self.newmodel.setEnabled(False)
         self.loadMasks.setEnabled(False)
         self.saveSet.setEnabled(False)
@@ -1485,14 +1484,11 @@ def update_plot(self):
             if self.color == 'rgb':
                 self.img.setImage(image, autoLevels=False, lut=None)
                 if self.nchan > 1:
-                    levels = np.array([
-                        self.saturation[0][self.currentZ],
-                        self.saturation[1][self.currentZ],
-                        self.saturation[2][self.currentZ]
-                    ])
-                    self.img.setLevels(levels)
-                else:
-                    self.img.setLevels(self.saturation[0][self.currentZ])
+                    levels = np.array([self.saturation[r][self.currentZ] for r in range(image.shape[-1])])
+                else: 
+                    levels = self.saturation[0][self.currentZ]
+                self.img.setLevels(levels)
+                
             elif self.color in rgb_list:
                 color_index = rgb_list.index(self.color)
                 if self.nchan > 1:
@@ -1505,6 +1501,7 @@ def update_plot(self):
             elif self.color == 'gray':
                 if self.nchan > 1:
                     # exclude channels with no data:
+                    # TODO: save this when computing saturation 
                     ranges = np.ptp(image, tuple(range(image.ndim-1)))
                     range_mask = ranges > 1e-5
                     image = image[..., range_mask]
@@ -1881,12 +1878,30 @@ 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])
-
-            if img_norm.shape[-1] == 1:
-                self.saturation.append(self.saturation[0])
-                self.saturation.append(self.saturation[0])
-
+        
+        elif len(self.saturation) == 0 or len(self.saturation[0]) != self.NZ:
+            self.saturation = []
+            for r in range(3):
+                self.saturation.append([])
+                for n in range(self.NZ):
+                    self.saturation[-1].append([0, 255])
+        
+        if img_norm.shape[-1] == 1:
+            if len(self.saturation) > 1:
+                self.saturation = [self.saturation[0]]
+            self.saturation.append(self.saturation[0])
+            self.saturation.append(self.saturation[0])
+        else:
+            self.saturation = [copy.deepcopy(self.saturation[r]) for r in range(img_norm.shape[-1])]
+            if len(self.saturation) == 2:
+                self.saturation.append([])
+                for n in range(self.NZ):
+                    self.saturation[-1].append([0, 255])
+        
+        self.sliders[0].setValue(self.saturation[0][self.currentZ])
+        self.sliders[1].setValue(self.saturation[1][self.currentZ])
+        self.sliders[2].setValue(self.saturation[2][self.currentZ])
+        # print(len(self.saturation), len(self.saturation[0]))
 
     def get_model_path(self, custom=False):
         if custom:

cellpose/gui/io.py

@@ -179,12 +179,16 @@ def _initialize_images(parent, image, load_3D=False):
 
     img_min = image.min()
     img_max = image.max()
-    parent.stack = parent.stack.astype(np.float32)
+    parent.stack = parent.stack.astype("float32")
     parent.stack -= img_min
     if img_max > img_min + 1e-3:
         parent.stack /= (img_max - img_min)
     parent.stack *= 255
 
+    if parent.stack.shape[-1] == 2:
+        parent.stack = np.concatenate((parent.stack, np.zeros((*parent.stack.shape[:-1], 1), dtype="float32")), axis=-1)
+    print(parent.stack.shape)
+
     if load_3D:
         parent.logger.info(": converted to float and normalized values to 0.0->255.0")
 
@@ -194,25 +198,12 @@ def _initialize_images(parent, image, load_3D=False):
     parent.imask = 0
     parent.Ly, parent.Lx = parent.stack.shape[-3:-1]
     parent.Ly0, parent.Lx0 = parent.stack.shape[-3:-1]
+    parent.nchan = parent.stack.shape[-1]
     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.Lyr, parent.Lxr = parent.stack.shape[-3:-1]
     parent.clear_all()
 
-    if not hasattr(parent, "stack_filtered") and parent.restore:
-        parent.logger.info(": no 'img_restore' found, applying current settings")
-        parent.compute_restore()
-
-    if parent.autobtn.isChecked():
-        if parent.restore is None or parent.restore != "filter":
-            parent.logger.info(": normalization checked: computing saturation levels (and optionally filtered image)")
-            parent.compute_saturation()
-    
+    parent.compute_saturation()    
     parent.compute_scale()
     parent.track_changes = []