fixtracks/fixtracks/detectionmerge.py

181 lines
7.2 KiB
Python

import logging
from PyQt6.QtWidgets import QWidget, QGridLayout, QVBoxLayout, QLabel, QPushButton, QComboBox, QSizePolicy, QSpinBox, QGraphicsView, QGraphicsScene, QGraphicsLineItem
from PyQt6.QtCore import QThreadPool, QLineF, QPointF, QRectF, Qt
from PyQt6.QtGui import QImage, QPixmap, QColor, QPen
from fixtracks.util import ImageReader
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):
def __init__(self, parent=None):
super().__init__(parent)
self._files = []
self.threadpool = QThreadPool()
self.right_imagereader = None
self.left_imagereader = None
layout = QVBoxLayout()
layout.addWidget(QLabel("Merge Detections!"))
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.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)
layout.addLayout(self.layout_controls())
splitter = QGridLayout()
splitter.addWidget(self.left_preview, 0, 0)
splitter.addWidget(self.right_preview, 0, 1)
splitter.setColumnStretch(0, 1)
splitter.setColumnStretch(1, 1)
layout.addLayout(splitter)
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(self.left_datacombo, 1, 1)
grd.addWidget(self.left_videocombo, 2, 1)
grd.addWidget(self.left_posspinner, 3, 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(self.right_datacombo, 1, 3)
grd.addWidget(self.right_videocombo, 2, 3)
grd.addWidget(self.right_posspinner, 3, 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.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(), 100)
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
img = self._toImage(frame)
self.right_preview.set_image(img)
self.right_posspinner.setMaximum(img.width() - 1)
self.right_posspinner.setValue(0)
def on_leftvideoSelection(self):
logging.debug("Video selection of the %s side", "left")
self.left_imagereader = ImageReader(self.left_videocombo.currentText(), 100)
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
img = self._toImage(frame)
self.left_preview.set_image(img)
self.left_posspinner.setMaximum(img.width() - 1)
self.left_posspinner.setValue(img.width() - 1)
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())