diff --git a/fixtracks/detectionmerge.py b/fixtracks/detectionmerge.py index bc45e5b..006a550 100644 --- a/fixtracks/detectionmerge.py +++ b/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!")) @@ -244,7 +251,7 @@ class MergeDetections(QWidget): if state: self._right_data = self.right_dataframereader.dataframe else: - self._right_data = None + self._right_data = None self.checkButtons() def on_leftmergelinemove(self): @@ -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, - self.left_posspinner.value(), - self.right_posspinner.value()) + 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.finished.connect(self.on_mergeDone) - self.threadpool.start(self.merger) + self._merger.signals.progress.connect(self.on_mergeProgress) + self._merger.signals.finished.connect(self.on_mergeDone) + 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() - self._progressDialog = None - self.merger = None + if state: + self._saveBtn.setEnabled(True) + # 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 - self._mergePreviewBtn.setEnabled(preview_enabled) + # preview_enabled = self.left_videocombo.currentIndex() > 0 and self.right_videocombo.currentIndex() > 0 + # 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.back.emit() \ No newline at end of file + self._merger = None + self._left_data = None + self._right_data = None + self.back.emit() diff --git a/fixtracks/util.py b/fixtracks/util.py index b18918f..b1d6c39 100644 --- a/fixtracks/util.py +++ b/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 - logging.info("Converting to numpy... %s", "Left camera") - lkeypoints, lquality, lbox = self.to_numpy(self._left_data) - 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.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) + 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._merged = self.to_dataframe(ldf, rdf, lkeypoints, rkeypoints, lboxes, rboxes, lquality, rquality, - lframes, rframes) + 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() \ No newline at end of file + main()