[merger] include merge result saving, some progress indicator

fixtracks/detectionmerge.py

@@ -5,6 +5,7 @@ from PyQt6.QtCore import QThreadPool, Qt, pyqtSignal, pyqtSlot
 from PyQt6.QtGui import QImage, QPixmap, QColor, QPen
 
 from fixtracks.util import ImageReader, DataFrameReader, Merger
+from PyQt6.QtWidgets import QFileDialog
 
 
 class VideoPreview(QWidget):
@@ -87,10 +88,15 @@ class MergeDetections(QWidget):
         self._mergePreviewBtn = QPushButton("Preview")
         self._mergePreviewBtn.clicked.connect(self.on_mergePreview)
         self._mergePreviewBtn.setEnabled(False)
+        self._mergePreviewBtn.setToolTip("Preview the merge results")
         self._mergeBtn = QPushButton("Merge!")
         self._mergeBtn.setEnabled(False)
         self._mergeBtn.setToolTip("Apply cutting and merge the data files into one")
         self._mergeBtn.clicked.connect(self.on_merge)
+        self._saveBtn = QPushButton("Save")
+        self._saveBtn.setToolTip("Save merge results")
+        self._saveBtn.clicked.connect(self.on_save)
+        self._saveBtn.setEnabled(False)
         self._backBtn = QPushButton("Back")
         self._backBtn.clicked.connect(self.on_back)
 
@@ -104,8 +110,9 @@ class MergeDetections(QWidget):
         btnBox.setAlignment(Qt.AlignmentFlag.AlignLeft)
         btnBox.addWidget(self._backBtn)
         btnBox.addItem(QSpacerItem(100, 10, QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed))
-        btnBox.addWidget(self._mergePreviewBtn)
         btnBox.addWidget(self._mergeBtn)
+        btnBox.addWidget(self._saveBtn)
+        btnBox.addWidget(self._mergePreviewBtn)
 
         layout = QVBoxLayout()
         layout.addWidget(QLabel("Merge Detections!"))
@@ -259,45 +266,69 @@ class MergeDetections(QWidget):
 
     def on_merge(self):
         logging.debug("detectionmerge: merge pressed")
-        self.merger = Merger(self._left_data, self._right_data,
+        if self._merger is not None:
+            self._merger = None
+        self._saveBtn.setEnabled(False)
+        self._merger = Merger(self._left_data, self._right_data,
                               self.left_posspinner.value(),
                               self.right_posspinner.value())
         self._progressDialog = QProgressDialog(parent=self)
         self._progressDialog.setAutoClose(True)
         self._progressDialog.setRange(0, 100)
-        self._progressDialog.setLabelText("Merging detections:")
+        self._progressDialog.setLabelText("Merging detections (will take a while, be patient):")
         self._progressDialog.setCancelButtonText("Cancel") 
         self._progressDialog.setWindowModality(Qt.WindowModality.WindowModal)
         self._progressDialog.canceled.connect(self.on_mergeCancelled)
         self._progressDialog.show()
-        self.merger.signals.progress.connect(self.on_mergeProgress)
+        self._merger.signals.progress.connect(self.on_mergeProgress)
-        self.merger.signals.finished.connect(self.on_mergeDone)
+        self._merger.signals.finished.connect(self.on_mergeDone)
-        self.threadpool.start(self.merger)
+        self.threadpool.start(self._merger)
 
     @pyqtSlot()
     def on_mergeCancelled(self):
-        self.merger.stop_request()
+        logging.info("Cancel Button pressed! Requesting stop of merger")
+        self._merger.stop_request()
+        self._saveBtn.setEnabled(False)
 
     @pyqtSlot(float)
     def on_mergeProgress(self, value):
-        print("mergeProgress", value)
+        logging.debug(f"mergeProgress: {value * 100}%")
         if self._progressDialog is not None:
             self._progressDialog.setValue(int(value * 100))
 
     @pyqtSlot(bool)
     def on_mergeDone(self, state):
         logging.debug("Merging stopped with status %s", state)
-        self._progressDialog.close()
+        if state:
-        self._progressDialog = None
+            self._saveBtn.setEnabled(True)
-        self.merger = None
+        # self._merger = None
 
     def checkButtons(self):
         merge_enabled = self._left_data is not None and self._right_data is not None
         logging.debug("CheckButtons: %s", str(merge_enabled))
         self._mergeBtn.setEnabled(merge_enabled)
-        preview_enabled = self.left_videocombo.currentIndex() > 0 and self.right_videocombo.currentIndex() > 0
+        # preview_enabled = self.left_videocombo.currentIndex() > 0 and self.right_videocombo.currentIndex() > 0
-        self._mergePreviewBtn.setEnabled(preview_enabled)
+        # self._mergePreviewBtn.setEnabled(preview_enabled)
+
+
+    def on_save(self):
+        logging.debug("Save merge results")
+        if self._merger is not None:
+            file_dialog = QFileDialog(self)
+            file_dialog.setAcceptMode(QFileDialog.AcceptMode.AcceptSave)
+            file_dialog.setNameFilter("Pickle Files (*.pkl)")
+            if file_dialog.exec():
+                file_path = file_dialog.selectedFiles()[0]
+                if not file_path.endswith(".pkl"):
+                    file_path += ".pkl"
+                self._merger.save(file_path)
+        else:
+            logging.debug("Saving failed! Merger is None!")
+
 
     def on_back(self):
         logging.debug("Back button pressed!")
+        self._merger = None
+        self._left_data = None
+        self._right_data = None
         self.back.emit()

fixtracks/util.py

@@ -241,7 +241,7 @@ class Merger(QRunnable):
             logging.error("Saving/pickling merged dataFrame is None!")
             return
         logging.info("Saving/pickling merged file to %s" % filename)
-        with open(filename, 'rb') as f:
+        with open(filename, 'wb') as f:
             pickle.dump(self._merged, f)
 
     @pyqtSlot()
@@ -251,27 +251,52 @@ class Merger(QRunnable):
     @pyqtSlot()
     def run(self):
         logging.info("Cutting left detections to limit %i", self._left_cut)
+        self.signals.progress.emit(0.0)
         self.signals.progress2.emit("Merging", self._mergeprogress, 0.)
         if not self.check_dataframe(self._left_data) or not self.check_dataframe(self._right_data):
             logging.error("Left or right dataframe structure does not match my expectations")
             return None
+        self.signals.progress.emit(0.05)
+
+        if not self._stopRequest:
             logging.info("Converting to numpy... %s", "Left camera")
             lkeypoints, lquality, lbox = self.to_numpy(self._left_data)
+            lframes = self._left_data.frame.values
+            self.signals.progress.emit(0.3)
+        else:
+            self.signals.finished(False)
+            return
+
+        if not self._stopRequest:
             logging.info("Converting to numpy... %s", "Right camera")
             rkeypoints, rquality, rbox = self.to_numpy(self._right_data)
-        lframes = self._left_data.frame.values
             rframes = self._right_data.frame.values
+            self.signals.progress.emit(0.6)
+        else:
+            self.signals.finished(False)
+            return
+
         logging.info("Filtering detections")
         left_easy, _ = self.sort_detections(lkeypoints, self._left_cut, left=True)
         right_easy, _ = self.sort_detections(rkeypoints, self._right_cut, left=False)
+        self.signals.progress.emit(0.7)
+
         logging.info("Merging and transformation")
         ldf, lkeypoints, lquality, lboxes, lframes = self.select_and_transform(self._left_data, lkeypoints, lbox,
                                                                                lquality, lframes, left_easy)
+        self.signals.progress.emit(0.8)
         rdf, rkeypoints, rquality, rboxes, rframes = self.select_and_transform(self._right_data, rkeypoints, rbox, 
                                                                                rquality, rframes, right_easy,
                                                                                self._left_cut, self._right_cut)
+        self.signals.progress.emit(0.9)
+
+        if not self._stopRequest:
             self._merged = self.to_dataframe(ldf, rdf, lkeypoints, rkeypoints, lboxes, rboxes, lquality, rquality,
                                              lframes, rframes)
+            self.signals.progress.emit(1.0)
+        else:
+            self.signals.finished(False)
+            return
 
         logging.info("Merging done!")
         self._signals.finished.emit(True and (not self._stopRequest))
@@ -285,6 +310,7 @@ class Merger(QRunnable):
         return self._result
 
 
+# TEST code
 def main():
     logging.info("Loading data left")
     left = pd.read_csv("../data/left_tracks.csv", sep=";", index_col=0)
@@ -297,5 +323,4 @@ def main():
 
 if __name__ == "__main__":
     logging.basicConfig(level=logging.INFO, force=True)
-
     main()