361 lines
15 KiB
Python
361 lines
15 KiB
Python
import logging
|
|
import pathlib
|
|
from PySide6.QtCore import Qt, QThreadPool, Signal, Slot, QRect, QSize, QPoint
|
|
from PySide6.QtWidgets import QWidget, QGridLayout, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QComboBox, QSizePolicy, QSpinBox, QGraphicsView, QGraphicsScene, QGraphicsLineItem, QSpacerItem, QProgressDialog, QFileDialog
|
|
from PySide6.QtGui import QImage, QPixmap, QColor, QPen, QPainter
|
|
|
|
from fixtracks.utils.reader import ImageReader, DataFrameReader
|
|
from fixtracks.utils.merger import Merger
|
|
|
|
|
|
class VideoPreview(QWidget):
|
|
def __init__(self, left=True, parent=None):
|
|
super().__init__(parent)
|
|
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._pixmapitem = None
|
|
self._left = left
|
|
layout = QVBoxLayout()
|
|
layout.addWidget(self._view)
|
|
self.setLayout(layout)
|
|
|
|
def setImage(self, img):
|
|
self._image = img
|
|
self._scene = QGraphicsScene()
|
|
self._pixmapitem = self._scene.addPixmap(QPixmap.fromImage(img))
|
|
self._view.setScene(self._scene)
|
|
self._view.fitInView(self._scene.sceneRect(), aspectRadioMode=Qt.AspectRatioMode.KeepAspectRatio)
|
|
image_rect = self._pixmapitem.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 setLineposition(self, position):
|
|
image_rect = self._pixmapitem.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)
|
|
|
|
@property
|
|
def image(self):
|
|
return self._image
|
|
|
|
|
|
class MergeDetections(QWidget):
|
|
back = Signal()
|
|
|
|
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._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)
|
|
|
|
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._mergeBtn)
|
|
btnBox.addWidget(self._saveBtn)
|
|
btnBox.addWidget(self._mergePreviewBtn)
|
|
|
|
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.setImage(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.setImage(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.setLineposition(self.left_posspinner.value())
|
|
|
|
def on_rightmergelinemove(self):
|
|
self.right_preview.setLineposition(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")
|
|
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 (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)
|
|
|
|
@Slot()
|
|
def on_mergeCancelled(self):
|
|
logging.info("Cancel Button pressed! Requesting stop of merger")
|
|
self._merger.stop_request()
|
|
self._saveBtn.setEnabled(False)
|
|
|
|
@Slot(float)
|
|
def on_mergeProgress(self, value):
|
|
logging.debug(f"mergeProgress: {value * 100}%")
|
|
if self._progressDialog is not None:
|
|
self._progressDialog.setValue(int(value * 100))
|
|
|
|
@Slot(bool)
|
|
def on_mergeDone(self, state):
|
|
logging.debug("Merging stopped with status %s", state)
|
|
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)
|
|
|
|
def merge_images(self):
|
|
limg = self.left_preview.image
|
|
left_cut = self.left_posspinner.value()
|
|
rimg = self.right_preview.image
|
|
right_cut = self.right_posspinner.value()
|
|
|
|
lrect = QRect(0, 0, left_cut, limg.height())
|
|
rrect = QRect(right_cut, 0, rimg.width() - right_cut, rimg.height())
|
|
lselection = limg.copy(lrect)
|
|
rselection = rimg.copy(rrect)
|
|
imgsize = QSize(lrect.width() + rrect.width(), rrect.height())
|
|
img = QImage(imgsize, limg.format())
|
|
|
|
painter = QPainter(img)
|
|
painter.drawImage(QPoint(0,0), lselection)
|
|
painter.drawImage(QPoint(lrect.width(), 0), rselection)
|
|
painter.end()
|
|
return img
|
|
|
|
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)
|
|
p = pathlib.Path(file_path)
|
|
img = self.merge_images()
|
|
img_name = p.with_suffix(".png")
|
|
img.save(str(img_name))
|
|
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()
|