fixtracks/fixtracks/widgets/mergepreview.py
2025-01-24 17:21:34 +01:00

156 lines
6.5 KiB
Python

import logging
import numpy as np
from PySide6.QtWidgets import QWidget, QVBoxLayout, QHBoxLayout, QLabel, QComboBox, QSizePolicy, QSpinBox, QGraphicsView, QGraphicsScene, QSpacerItem, QSlider
from PySide6.QtCore import Qt
from PySide6.QtGui import QPixmap, QBrush, QColor
class MergePreview(QWidget):
def __init__(self, parent=None):
super().__init__(parent)
self._dataframe = None
self._image = None
self._pixmapitem = None
self._frames = None
self._keypoints = None
self._scene = None
self._view = QGraphicsView()
self._view.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
self._brush = QBrush(QColor.fromString("red"))
self._brush2 = QBrush(QColor.fromString("white"))
self._startSlider = QSlider(orientation=Qt.Orientation.Horizontal)
self._startSlider.setTickPosition(QSlider.TickPosition.TicksBelow)
self._startSlider.setRange(0, 100)
self._startSlider.setTickInterval(self._startSlider.maximum() // 10)
self._startSlider.sliderReleased.connect(self._on_startSliderReleased)
self._startSlider.valueChanged.connect(self._on_sliderChanged)
self._stopSlider = QSlider(orientation=Qt.Orientation.Horizontal)
self._stopSlider.setTickPosition(QSlider.TickPosition.TicksBelow)
self._stopSlider.setRange(0, 100)
self._stopSlider.setTickInterval(self._stopSlider.maximum() // 10)
self._stopSlider.sliderReleased.connect(self._on_stopSliderReleased)
self._stopSlider.valueChanged.connect(self._on_sliderChanged)
self._keypointCombo = QComboBox()
self._keypointCombo.currentIndexChanged.connect(self._on_keypointSelection)
self._numpointsSpinner = QSpinBox()
self._numpointsSpinner.setRange(500, 50000)
self._numpointsSpinner.setValue(25000)
self._numpointsSpinner.setSingleStep(500)
self._numpointsSpinner.valueChanged.connect(self.update)
settingsbox = QHBoxLayout()
settingsbox.addWidget(QLabel("Select keypoint"))
settingsbox.addWidget(self._keypointCombo)
settingsbox.addWidget(QLabel("Number of displayed points"))
settingsbox.addWidget(self._numpointsSpinner)
startsliderbox = QHBoxLayout()
startsliderbox.addWidget(QLabel("Start position"))
startsliderbox.addWidget(self._startSlider)
stopsliderbox = QHBoxLayout()
stopsliderbox.addWidget(QLabel("Stop position"))
stopsliderbox.addWidget(self._stopSlider)
self._startLabel = QLabel("0")
self._stopLabel = QLabel("100")
messagebox = QHBoxLayout()
messagebox.addWidget(QLabel("Current segment from frame:"))
messagebox.addWidget(self._startLabel)
messagebox.addWidget(QLabel("to"))
messagebox.addWidget(self._stopLabel)
messagebox.addSpacerItem(QSpacerItem(100, 10, QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed))
layout = QVBoxLayout()
layout.addWidget(self._view)
layout.addLayout(settingsbox)
layout.addLayout(startsliderbox)
layout.addLayout(stopsliderbox)
layout.addLayout(messagebox)
self.setLayout(layout)
def resetSlider(self, framecount=100):
self._startSlider.setRange(0, framecount)
self._stopSlider.setRange(0, framecount)
self._startSlider.setSliderPosition(0)
self._stopSlider.setSliderPosition(framecount)
self._startSlider.setTickInterval(self._startSlider.maximum() // 10)
self._stopSlider.setTickInterval(self._stopSlider.maximum() // 10)
def resetKeypointcombo(self, numkeypoints):
self._keypointCombo.clear()
self._keypointCombo.addItems([f"keypoint {i}" for i in range(numkeypoints)])
def setDataframe(self, newdf):
self._dataframe = newdf
self._frames = self._dataframe.frame.values
self._keypoints = self._dataframe.keypoints.values
self.resetSlider(np.max(self._dataframe.frame))
self.resetKeypointcombo(len(self._keypoints[0]))
self.update()
def _on_keypointSelection(self):
self.update()
def _on_startSliderReleased(self):
value = self._startSlider.value()
if value > self._stopSlider.value():
self._stopSlider.setValue(value)
self.update()
def _on_stopSliderReleased(self):
value = self._stopSlider.value()
if value < self._startSlider.value():
self._startSlider.setValue(value)
self.update()
def _on_sliderChanged(self, value):
self._startLabel.setText(str(self._startSlider.value()))
self._stopLabel.setText(str(self._stopSlider.value()))
def setImage(self, img):
self._image = img
self.update()
def _drawDetections(self):
start_frame = self._startSlider.value()
stop_frame = self._stopSlider.value()
selection = self._keypoints[self._frames[(self._frames >= start_frame) & (self._frames < stop_frame)]]
image_rect = self._pixmapitem.boundingRect()
keypoint = self._keypointCombo.currentIndex()
num_points = self._numpointsSpinner.value()
step = 1 if len(selection) < num_points else int(len(selection) // num_points)
for coords in selection[::step]:
x = coords[keypoint, 0]
y = coords[keypoint, 1]
self._scene.addEllipse(image_rect.left() + x, image_rect.top() + y, 10, 10, brush=self._brush)
def update(self):
if self._image is not None:
self._scene = QGraphicsScene()
self._pixmapitem = self._scene.addPixmap(QPixmap.fromImage(self._image))
self._view.setScene(self._scene)
self._view.fitInView(self._scene.sceneRect(), aspectRadioMode=Qt.AspectRatioMode.KeepAspectRatio)
self._view.show()
if self._dataframe is not None:
self._scene = QGraphicsScene()
self._pixmapitem = self._scene.addPixmap(QPixmap.fromImage(self._image))
self._view.setScene(self._scene)
self._view.fitInView(self._scene.sceneRect(), aspectRadioMode=Qt.AspectRatioMode.KeepAspectRatio)
self._view.show()
self._drawDetections()
def fit_image_to_view(self):
"""Scale the image to fit the QGraphicsView."""
if self._pixmapitem is not None:
self._view.fitInView(self._pixmapitem, Qt.KeepAspectRatio)
def resizeEvent(self, event):
"""Handle window resizing to fit the image."""
super().resizeEvent(event)
self.fit_image_to_view()