diff --git a/fixtracks/widgets/tracks.py b/fixtracks/widgets/tracks.py index 9d8f3e2..66f78ce 100644 --- a/fixtracks/widgets/tracks.py +++ b/fixtracks/widgets/tracks.py @@ -11,7 +11,7 @@ from PySide6.QtWidgets import QSpinBox, QSpacerItem, QProgressBar, QSplitter, QG from fixtracks.utils.reader import PickleLoader from fixtracks.utils.writer import PickleWriter -from fixtracks.widgets.detectionview import DetectionView +from fixtracks.widgets.detectionview import DetectionView, DetectionData from fixtracks.widgets.detectiontimeline import DetectionTimeline from fixtracks.widgets.skeleton import SkeletonWidget @@ -385,11 +385,11 @@ class FixTracks(QWidget): btnBox.addWidget(self._tasklabel) btnBox.addWidget(self._progress_bar) btnBox.addWidget(self._saveBtn) - + cntrlBox = QHBoxLayout() - cntrlBox.addItem(QSpacerItem(100, 100, QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)) - cntrlBox.addWidget(self._controls_widget, stretch=1, alignment=Qt.AlignmentFlag.AlignCenter) - cntrlBox.addWidget(self._skeleton, stretch=1, alignment=Qt.AlignmentFlag.AlignCenter) + cntrlBox.addItem(QSpacerItem(200, 100, QSizePolicy.Policy.Fixed, QSizePolicy.Policy.Expanding)) + cntrlBox.addWidget(self._controls_widget) + cntrlBox.addWidget(self._skeleton) vbox = QVBoxLayout() vbox.addLayout(timelinebox) vbox.addLayout(cntrlBox) @@ -570,15 +570,23 @@ class FixTracks(QWidget): def on_detectionsSelected(self, detections): logging.debug("Tracks: Detections selected") - tracks = np.zeros(len(detections)) + tracks = np.zeros(len(detections), dtype=int) ids = np.zeros_like(tracks) - coordinates = np.zeros((len(detections), )) + frames = np.zeros_like(tracks) + coordinates = None + if len(detections) > 0: + c = detections[0].data(DetectionData.COORDINATES.value) + coordinates = np.zeros((len(detections), c.shape[0], c.shape[1])) + for i, d in enumerate(detections): - tracks[i] = d.data(0) - ids[i] = d.data(1) - + tracks[i] = d.data(DetectionData.TRACK_ID.value) + ids[i] = d.data(DetectionData.ID.value) + frames[i] = d.data(DetectionData.FRAME.value) + coordinates[i, :, :] = d.data(DetectionData.COORDINATES.value) self._data.setUserSelection(ids) self._controls_widget.setSelectedTracks(tracks) + self._skeleton.clear() + self._skeleton.addSkeletons(coordinates, ids, QBrush(QColor(10, 255, 65, 255))) self.update() def moveWindow(self, stepsize):