fixtracks/fixtracks/detectionmerge.py
2024-12-30 13:46:16 +01:00

303 lines
12 KiB
Python

import logging
from PyQt6.QtWidgets import QWidget, QGridLayout, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QComboBox, QSizePolicy, QSpinBox, QGraphicsView, QGraphicsScene, QGraphicsLineItem, QSpacerItem, QProgressDialog, QProgressBar
from PyQt6.QtCore import QThreadPool, Qt, pyqtSignal, pyqtSlot
from PyQt6.QtGui import QImage, QPixmap, QColor, QPen
from fixtracks.util import ImageReader, DataFrameReader, Merger
class VideoPreview(QWidget):
def __init__(self, left=True, parent=None):
super().__init__(None)
self._image = None
self._line_pen = QPen(QColor.fromString("red"), 4)
self._view = QGraphicsView()
self._view.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
self._scene = None
self._pixmap = None
self._left = left
layout = QVBoxLayout()
layout.addWidget(self._view)
self.setLayout(layout)
def set_image(self, img):
self._scene = QGraphicsScene()
self._pixmap = self._scene.addPixmap(QPixmap.fromImage(img))
self._view.setScene(self._scene)
self._view.fitInView(self._scene.sceneRect(), mode=Qt.AspectRatioMode.KeepAspectRatio)
image_rect = self._pixmap.boundingRect()
start_x = image_rect.right() if self._left else image_rect.left()
start_y = image_rect.top()
end_y = image_rect.bottom()
self._scene.addLine(start_x, start_y, start_x, end_y, self._line_pen)
self._view.show()
def set_lineposition(self, position):
image_rect = self._pixmap.boundingRect()
for i in self._scene.items():
if isinstance(i, QGraphicsLineItem):
self._scene.removeItem(i)
start_x = image_rect.left() + position
start_y = image_rect.top()
end_y = image_rect.bottom()
self._scene.addLine(start_x, start_y, start_x, end_y, self._line_pen)
class MergeDetections(QWidget):
back = pyqtSignal()
def __init__(self, parent=None):
super().__init__(parent)
self._files = []
self.threadpool = QThreadPool()
self.right_imagereader = None
self.left_imagereader = None
self.right_dataframereader = None
self.left_dataframereader = None
self._left_data = None
self._right_data = None
self._progressDialog = None
self._merger = None
self.left_datacombo = QComboBox()
self.left_videocombo = QComboBox()
self.left_datacombo.addItems(self._files)
self.left_videocombo.addItems(self._files)
self.left_posspinner = QSpinBox()
self.left_posspinner.valueChanged.connect(self.on_leftmergelinemove)
self.left_preview = VideoPreview()
self.left_framespinner = QSpinBox()
self.left_framespinner.setMaximum(10000)
self.left_framespinner.setValue(100)
self.right_datacombo = QComboBox()
self.right_videocombo = QComboBox()
self.right_datacombo.addItems(self._files)
self.right_videocombo.addItems(self._files)
self.right_posspinner = QSpinBox()
self.right_posspinner.valueChanged.connect(self.on_rightmergelinemove)
self.right_preview = VideoPreview(left=False)
self.right_framespinner = QSpinBox()
self.right_framespinner.setMaximum(10000)
self.right_framespinner.setValue(100)
self._mergePreviewBtn = QPushButton("Preview")
self._mergePreviewBtn.clicked.connect(self.on_mergePreview)
self._mergePreviewBtn.setEnabled(False)
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._backBtn = QPushButton("Back")
self._backBtn.clicked.connect(self.on_back)
grid = QGridLayout()
grid.addWidget(self.left_preview, 0, 0)
grid.addWidget(self.right_preview, 0, 1)
grid.setColumnStretch(0, 1)
grid.setColumnStretch(1, 1)
btnBox = QHBoxLayout()
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)
layout = QVBoxLayout()
layout.addWidget(QLabel("Merge Detections!"))
layout.addLayout(self.layout_controls())
layout.addLayout(grid)
layout.addLayout(btnBox)
self.setLayout(layout)
def layout_controls(self):
grd = QGridLayout()
grd.addWidget(QLabel("Left:"), 0, 0)
grd.addWidget(QLabel("Data"), 1, 0)
grd.addWidget(QLabel("Video"), 2, 0)
grd.addWidget(QLabel("Mergeline"), 3, 0)
grd.addWidget(QLabel("Seek frame"), 4, 0)
grd.addWidget(self.left_datacombo, 1, 1)
grd.addWidget(self.left_videocombo, 2, 1)
grd.addWidget(self.left_posspinner, 3, 1)
grd.addWidget(self.left_framespinner, 4, 1)
grd.addWidget(QLabel("Right"), 0, 2)
grd.addWidget(QLabel("Data"), 1, 2)
grd.addWidget(QLabel("Video"), 2, 2)
grd.addWidget(QLabel("Mergeline"), 3, 2)
grd.addWidget(QLabel("Seek frame"), 4, 2)
grd.addWidget(self.right_datacombo, 1, 3)
grd.addWidget(self.right_videocombo, 2, 3)
grd.addWidget(self.right_posspinner, 3, 3)
grd.addWidget(self.right_framespinner, 4, 3)
grd.setColumnStretch(0, 0)
grd.setColumnStretch(2, 0)
grd.setColumnStretch(1, 50)
grd.setColumnStretch(3, 50)
grd.setColumnMinimumWidth(0, 50)
grd.setColumnMinimumWidth(2, 50)
return grd
@property
def fileList(self):
return self._files
@fileList.setter
def fileList(self, file_list):
logging.debug("MergeDetections.fileList: set new file list")
logging.debug("MergeDetections.fileList: setting video combo boxes")
videoformats = [".mp4", ".avi"]
self._files = [str(f) for f in file_list if f.suffix in videoformats]
self.right_videocombo.addItem("Please select")
self.left_videocombo.addItem("Please select")
self.right_videocombo.addItems(self.fileList)
self.left_videocombo.addItems(self.fileList)
self.left_videocombo.setCurrentIndex(0)
self.right_videocombo.setCurrentIndex(0)
logging.debug("MergeDetections.fileList: setting data combo boxes")
dataformats = [".csv"]
self._files = [str(f) for f in file_list if f.suffix in dataformats]
self.right_datacombo.addItem("Please select")
self.left_datacombo.addItem("Please select")
self.right_datacombo.addItems(self.fileList)
self.left_datacombo.addItems(self.fileList)
self.left_datacombo.setCurrentIndex(0)
self.right_datacombo.setCurrentIndex(0)
self.left_datacombo.currentIndexChanged.connect(self.on_leftDataSelection)
self.right_datacombo.currentIndexChanged.connect(self.on_rightDataSelection)
self.right_videocombo.currentIndexChanged.connect(self.on_rightvideoSelection)
self.left_videocombo.currentIndexChanged.connect(self.on_leftVideoSelection)
def _toImage(self, frame):
height, width, _ = frame.shape
bytesPerLine = 3 * width
img = QImage(frame.data, width, height, bytesPerLine, QImage.Format.Format_BGR888).rgbSwapped()
return img
def on_rightvideoSelection(self):
logging.debug("Video selection of the %s side", "right")
self.right_imagereader = ImageReader(self.right_videocombo.currentText(), self.right_framespinner.value())
self.right_imagereader.signals.finished.connect(self.right_imgreaderDone)
self.threadpool.start(self.right_imagereader)
def right_imgreaderDone(self, state):
logging.debug("Right image reader done with state %s", str(state))
frame = self.right_imagereader.frame
if frame is not None:
img = self._toImage(frame)
self.right_preview.set_image(img)
self.right_posspinner.setMaximum(img.width() - 1)
self.right_posspinner.setValue(0)
self.checkButtons()
else:
logging.error("Reading frame failed!")
def on_leftVideoSelection(self):
logging.debug("Video selection of the %s side", "left")
self.left_imagereader = ImageReader(self.left_videocombo.currentText(), self.left_framespinner.value())
self.left_imagereader.signals.finished.connect(self.left_imgreaderDone)
self.threadpool.start(self.left_imagereader)
def left_imgreaderDone(self, state):
logging.debug("Left image reader done with state %s", str(state))
frame = self.left_imagereader.frame
if frame is not None:
img = self._toImage(frame)
self.left_preview.set_image(img)
self.left_posspinner.setMaximum(img.width() - 1)
self.left_posspinner.setValue(img.width() - 1)
self.checkButtons()
else:
logging.error("Reading frame failed!")
def on_leftDataSelection(self):
logging.debug("Data selection of the %s side", "left")
self.left_dataframereader = DataFrameReader(self.left_datacombo.currentText())
self.left_dataframereader.signals.finished.connect(self.left_dfreaderDone)
self.threadpool.start(self.left_dataframereader)
def left_dfreaderDone(self, state):
logging.debug("Left data reader done with state %s", str(state))
if state:
self._left_data = self.left_dataframereader.dataframe
else:
self._left_data = None
self.checkButtons()
def on_rightDataSelection(self):
logging.debug("Data selection of the %s side", "right")
self.right_dataframereader = DataFrameReader(self.right_datacombo.currentText())
self.right_dataframereader.signals.finished.connect(self.right_dfreaderDone)
self.threadpool.start(self.right_dataframereader)
def right_dfreaderDone(self, state):
logging.debug("Right data reader done with state %s", str(state))
if state:
self._right_data = self.right_dataframereader.dataframe
else:
self._right_data = None
self.checkButtons()
def on_leftmergelinemove(self):
self.left_preview.set_lineposition(self.left_posspinner.value())
def on_rightmergelinemove(self):
self.right_preview.set_lineposition(self.right_posspinner.value())
def on_mergePreview(self):
logging.debug("detectionmerge: mergePreview pressed")
print("Sorry, not implemented yet!")
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())
self._progressDialog = QProgressDialog(parent=self)
self._progressDialog.setAutoClose(True)
self._progressDialog.setRange(0, 100)
self._progressDialog.setLabelText("Merging detections:")
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)
@pyqtSlot()
def on_mergeCancelled(self):
self.merger.stop_request()
@pyqtSlot(float)
def on_mergeProgress(self, value):
print("mergeProgress", value)
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
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)
def on_back(self):
logging.debug("Back button pressed!")
self.back.emit()