fix: segm multiple frames in 2D from 4D data in gui, see #354

cellacdc/apps.py

@@ -5710,7 +5710,7 @@ def __init__(self, mainWindow):
         self.plotMarkers()
 
     def getImage(self):
-        img = self.mainWindow.getDisplayedCellsImg()
+        img = self.mainWindow.getDisplayedImg1()
         self.img = img/img.max()
         self.imgRGB = (skimage.color.gray2rgb(self.img)*255).astype(np.uint8)
 

cellacdc/gui.py

@@ -4497,7 +4497,7 @@ def gui_mousePressEventImg2(self, event):
             # If automatic bud separation was not successfull call manual one
             if not success:
                 posData.disableAutoActivateViewerWindow = True
-                img = self.getDisplayedCellsImg()
+                img = self.getDisplayedImg1()
                 col = 'manual_separate_draw_mode'
                 drawMode = self.df_settings.at[col, 'value']
                 manualSep = apps.manualSeparateGui(
@@ -6774,7 +6774,7 @@ def gui_mousePressEventImg1(self, event):
             if eraseOnlyOneID:
                 mask[lab_2D!=self.erasedID] = False
 
-            self.getDisplayedCellsImg()
+            self.getDisplayedImg1()
             self.setTempImg1Eraser(mask, init=True)
             self.applyEraserMask(mask)
 
@@ -6886,7 +6886,7 @@ def gui_mousePressEventImg1(self, event):
             self.autoCont_y0 = ydata
             self.xxA_autoCont, self.yyA_autoCont = [], []
             self.curvAnchors.addPoints([x], [y])
-            img = self.getDisplayedCellsImg()
+            img = self.getDisplayedImg1()
             self.autoContObjMask = np.zeros(img.shape, np.uint8)
             self.isRightClickDragImg1 = True
 
@@ -6950,7 +6950,7 @@ def gui_mousePressEventImg1(self, event):
 
             # NOTE: flood is on mousedrag or release
             tol = self.wandToleranceSlider.value()
-            self.flood_img = myutils.to_uint8(self.getDisplayedCellsImg())
+            self.flood_img = myutils.to_uint8(self.getDisplayedImg1())
             flood_mask = skimage.segmentation.flood(
                 self.flood_img, (ydata, xdata), tolerance=tol
             )
@@ -8064,7 +8064,7 @@ def drawAutoContour(self, y2, x2):
         y1, x1 = self.autoCont_y0, self.autoCont_x0
         Dy = abs(y2-y1)
         Dx = abs(x2-x1)
-        edge = self.getDisplayedCellsImg()
+        edge = self.getDisplayedImg1()
         if Dy != 0 or Dx != 0:
             # NOTE: numIter takes care of any lag in mouseMoveEvent
             numIter = int(round(max((Dy, Dx))))
@@ -12778,7 +12778,7 @@ def segmWorkerFinished(self, lab, exec_time):
         self.checkIfAutoSegm()
 
     # @exec_time
-    def getDisplayedCellsImg(self):
+    def getDisplayedImg1(self):
         return self.img1.image
 
     def getDisplayedZstack(self):
@@ -12858,7 +12858,7 @@ def autoAssignBud_YeastMate(self):
 
         self.models[idx] = model
 
-        img = self.getDisplayedCellsImg()
+        img = self.getDisplayedImg1()
 
         posData.cca_df = model.predictCcaState(img, posData.lab)
         self.store_data()
@@ -13979,12 +13979,12 @@ def init_segmInfo_df(self):
             if posData.SizeZ > 1 and posData.segmInfo_df is not None:
                 if 'z_slice_used_gui' not in posData.segmInfo_df.columns:
                     posData.segmInfo_df['z_slice_used_gui'] = (
-                                    posData.segmInfo_df['z_slice_used_dataPrep']
-                                    )
+                        posData.segmInfo_df['z_slice_used_dataPrep']
+                    )
                 if 'which_z_proj_gui' not in posData.segmInfo_df.columns:
                     posData.segmInfo_df['which_z_proj_gui'] = (
-                                    posData.segmInfo_df['which_z_proj']
-                                    )
+                        posData.segmInfo_df['which_z_proj']
+                    )
                 posData.segmInfo_df['resegmented_in_gui'] = False
                 posData.segmInfo_df.to_csv(posData.segmInfo_df_csv_path)
 

cellacdc/workers.py

@@ -481,7 +481,7 @@ def run(self):
                 startZ, stopZ = self.z_range
                 img = img[startZ:stopZ+1]
         else:
-            img = self.mainWin.getDisplayedCellsImg()
+            img = self.mainWin.getDisplayedImg1()
 
         posData = self.mainWin.data[self.mainWin.pos_i]
         lab = np.zeros_like(posData.segm_data[0])
@@ -534,15 +534,44 @@ def _check_extend_segm_data(self, segm_data, stop_frame_num):
         extended_shape = (stop_frame_num, *segm_data.shape[1:])
         extended_segm_data = np.zeros(extended_shape, dtype=segm_data.dtype)
         extended_segm_data[:len(segm_data)] = segm_data
+        if len(extended_shape) == 4:
+            return extended_segm_data
+        if self.posData.SizeZ == 1:
+            return extended_segm_data
+        else:
+            num_added_frames = len(extended_segm_data) - len(segm_data)
+            half_z = int(self.posData.SizeZ/2)
+            # 2D segm on 3D over time data --> fix segmInfo
+            segmInfo_extended = pd.DataFrame({
+                'filename': [self.posData.filename]*num_added_frames,
+                'frame_i': list(range(len(segm_data), len(extended_segm_data))),
+                'z_slice_used_gui': [half_z]*num_added_frames,
+                'which_z_proj_gui': ['single z-slice']*num_added_frames
+            }).set_index(['filename', 'frame_i'])
+            segmInfo_df = pd.concat([self.posData.segmInfo_df, segmInfo_extended])
+            self.posData.segmInfo_df = segmInfo_df
+            self.posData.segmInfo_df.to_csv(self.posData.segmInfo_df_csv_path)
         return extended_segm_data
+
+    def getImageData(self):
+        imgData = self.posData.img_data[self.startFrameNum-1:self.stopFrameNum]
+        if self.posData.SizeZ == 1:
+            return imgData
+        if self.posData.segm_data.ndim == 4:
+            return imgData
+        # 2D segm on 3D over time data --> index z-slices
+        filename = self.posData.filename
+        zz = self.posData.segmInfo_df.loc[filename, 'z_slice_used_gui']
+        imgData = self.posData.img_data[:, zz.to_list()]
+        return imgData
 
     @worker_exception_handler
     def run(self):
         t0 = time.perf_counter()
-        img_data = self.posData.img_data[self.startFrameNum-1:self.stopFrameNum]
         self.posData.segm_data = self._check_extend_segm_data(
             self.posData.segm_data, self.stopFrameNum
         )
+        img_data = self.getImageData()
         for i, img in enumerate(img_data):
             frame_i = i+self.startFrameNum-1
             if self.secondChannelData is not None: