193 lines
7.3 KiB
Python
193 lines
7.3 KiB
Python
import logging
|
|
import pathlib
|
|
import numpy as np
|
|
|
|
from PySide6.QtCore import Qt, QThreadPool, Signal
|
|
from PySide6.QtGui import QImage, QStandardItemModel, QStandardItem
|
|
from PySide6.QtWidgets import QWidget, QVBoxLayout, QHBoxLayout, QLabel, QPushButton, QSizePolicy
|
|
from PySide6.QtWidgets import QSpinBox, QSpacerItem, QFileDialog, QProgressBar, QTableView, QSplitter
|
|
|
|
from fixtracks.utils.reader import PickleLoader
|
|
from fixtracks.widgets.detectionview import DetectionView
|
|
from fixtracks.widgets.timeline import Timeline
|
|
|
|
class FixTracks(QWidget):
|
|
back = Signal()
|
|
|
|
def __init__(self, parent=None):
|
|
super().__init__(parent)
|
|
self._files = []
|
|
self._threadpool = QThreadPool()
|
|
self._reader = None
|
|
self._image = None
|
|
self._dataframe = None
|
|
self._unassignedmodel = None
|
|
self._leftmodel = None
|
|
self._rightmodel = None
|
|
|
|
self._detectionView = DetectionView()
|
|
self._progress_bar = QProgressBar(self)
|
|
self._progress_bar.setMaximumHeight(20)
|
|
# self._progress_bar.setRange(0, 0) # Set the progress bar to be indeterminate
|
|
self._progress_bar.setValue(0)
|
|
self._tasklabel = QLabel()
|
|
|
|
self._timeline = Timeline()
|
|
self._timeline.signals.changed.connect(self.on_windowChanged)
|
|
self._windowspinner = QSpinBox()
|
|
self._windowspinner.setRange(100, 10000)
|
|
self._windowspinner.setSingleStep(100)
|
|
self._windowspinner.valueChanged.connect(self.on_windowSizeChanged)
|
|
|
|
timelinebox = QHBoxLayout()
|
|
timelinebox.addWidget(self._timeline)
|
|
timelinebox.addWidget(QLabel("Window"))
|
|
timelinebox.addWidget(self._windowspinner)
|
|
|
|
self._left_table = QTableView()
|
|
assign1 = QPushButton("<<")
|
|
assign1.clicked.connect(self.on_assignLeft)
|
|
assign2 = QPushButton(">>")
|
|
assign2.clicked.connect(self.on_assignRight)
|
|
self._unassigned_table = QTableView()
|
|
self._right_table = QTableView()
|
|
tablebox = QHBoxLayout()
|
|
tablebox.addWidget(self._left_table)
|
|
tablebox.addWidget(assign1)
|
|
tablebox.addWidget(self._unassigned_table)
|
|
tablebox.addWidget(assign2)
|
|
tablebox.addWidget(self._right_table)
|
|
|
|
self._openBtn = QPushButton("Open")
|
|
self._openBtn.setEnabled(True)
|
|
self._openBtn.clicked.connect(self._on_open)
|
|
self._saveBtn = QPushButton("Save")
|
|
self._saveBtn.setEnabled(False)
|
|
self._saveBtn.clicked.connect(self.on_save)
|
|
self._backBtn = QPushButton("Back")
|
|
self._backBtn.clicked.connect(self.on_back)
|
|
|
|
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._tasklabel)
|
|
btnBox.addWidget(self._progress_bar)
|
|
btnBox.addWidget(self._openBtn)
|
|
btnBox.addWidget(self._saveBtn)
|
|
|
|
vbox = QVBoxLayout()
|
|
vbox.addLayout(timelinebox)
|
|
vbox.addLayout(tablebox)
|
|
vbox.addLayout(btnBox)
|
|
container = QWidget()
|
|
container.setLayout(vbox)
|
|
|
|
splitter = QSplitter(Qt.Orientation.Horizontal)
|
|
splitter.addWidget(self._detectionView)
|
|
splitter.addWidget(container)
|
|
splitter.setStretchFactor(0, 3)
|
|
splitter.setStretchFactor(1, 1)
|
|
layout = QHBoxLayout()
|
|
layout.addWidget(splitter)
|
|
self.setLayout(layout)
|
|
|
|
def _on_open(self):
|
|
infile = None
|
|
imgfile = None
|
|
|
|
self._tasklabel.setText( "Select merged image")
|
|
file_dialog = QFileDialog(self, "Select merged image")
|
|
file_dialog.setFileMode(QFileDialog.ExistingFile)
|
|
file_dialog.setNameFilters([
|
|
"Image Files (*.png *.jpg *.jpeg)",
|
|
"All Files (*)"
|
|
])
|
|
if file_dialog.exec():
|
|
imgfile = file_dialog.selectedFiles()[0]
|
|
if imgfile is not None:
|
|
img = QImage(imgfile)
|
|
self._detectionView.setImage(img)
|
|
self._tasklabel.setText( "Open data")
|
|
file_dialog = QFileDialog(self, "Select pickled DataFrame", "", "Pandas DataFrame (*.pkl)")
|
|
if file_dialog.exec():
|
|
infile = file_dialog.selectedFiles()[0]
|
|
if infile is not None:
|
|
self._progress_bar.setRange(0,0)
|
|
self._reader = PickleLoader(infile)
|
|
self._reader.signals.finished.connect(self._on_dataOpenend)
|
|
self._threadpool.start(self._reader)
|
|
|
|
|
|
def populateTables(self):
|
|
left_trackid = 1
|
|
right_trackid = 2
|
|
start_frame = self._timeline.sliderPosition - self._windowspinner.value() // 2
|
|
stop_frame = self._timeline.sliderPosition + self._windowspinner.value() // 2
|
|
df = self._dataframe[(self._dataframe.frame >= start_frame) & (self._dataframe.frame < stop_frame)]
|
|
|
|
assigned_left = df[(df.track == left_trackid)]
|
|
assigned_right = df[(df.track == right_trackid)]
|
|
unassigned = df[(df.track != left_trackid) & (df.track != right_trackid)]
|
|
print(len(assigned_left), len(assigned_right), len(unassigned))
|
|
columns = ["frame", "track id"]
|
|
self._unassignedmodel = QStandardItemModel(len(unassigned), 2)
|
|
self._unassignedmodel.setHorizontalHeaderLabels(columns)
|
|
self._leftmodel = QStandardItemModel(len(assigned_left), 2)
|
|
self._leftmodel.setHorizontalHeaderLabels(columns)
|
|
self._rightmodel = QStandardItemModel(len(assigned_right), 2)
|
|
self._rightmodel.setHorizontalHeaderLabels(columns)
|
|
|
|
# Populate the models with data
|
|
for i in range(len(unassigned)):
|
|
row = unassigned.iloc[i]
|
|
if i == 0: print(row)
|
|
for j in range(len(columns)):
|
|
item = QStandardItem(f"{i, j}")
|
|
self._unassignedmodel.setItem(i, j, item)
|
|
self._unassigned_table.setModel(self._unassignedmodel)
|
|
|
|
for i in range(len(assigned_left)):
|
|
row = assigned_left.iloc[i]
|
|
for j in range(len(columns)):
|
|
item = QStandardItem(f"{i, j}")
|
|
self._leftmodel.setItem(i, j, item)
|
|
self._left_table.setModel(self._leftmodel)
|
|
|
|
for i in range(len(assigned_right)):
|
|
row = assigned_right.iloc[i]
|
|
for j in range(len(columns)):
|
|
item = QStandardItem(f"{i, j}")
|
|
self._rightmodel.setItem(i, j, item)
|
|
self._right_table.setModel(self._rightmodel)
|
|
|
|
|
|
def _on_dataOpenend(self, state):
|
|
logging.info("Finished loading data with state %s", state)
|
|
self._tasklabel.setText("")
|
|
self._progress_bar.setRange(0, 100)
|
|
self._progress_bar.setValue(0)
|
|
if state and self._reader is not None:
|
|
self._dataframe = self._reader.data
|
|
self._timeline.setRange(np.max(self._dataframe.frame.values), self._windowspinner.value())
|
|
self.populateTables()
|
|
|
|
def on_save(self):
|
|
logging.debug("Save fixtracks results")
|
|
|
|
def on_back(self):
|
|
logging.debug("Back button pressed!")
|
|
self.back.emit()
|
|
|
|
def on_assignLeft(self):
|
|
pass
|
|
|
|
def on_assignRight(self):
|
|
pass
|
|
|
|
def on_windowChanged(self, value):
|
|
self.populateTables()
|
|
|
|
def on_windowSizeChanged(self, value):
|
|
self._timeline.setWindowWidth(value)
|