diff --git a/fixtracks/widgets/skeleton.py b/fixtracks/widgets/skeleton.py index 4e8038e..1f961e7 100644 --- a/fixtracks/widgets/skeleton.py +++ b/fixtracks/widgets/skeleton.py @@ -1,24 +1,36 @@ import logging import numpy as np -from PySide6.QtWidgets import QWidget, QVBoxLayout, QSizePolicy, QGraphicsView, QGraphicsItem, QGraphicsScene, QGraphicsEllipseItem, QGraphicsRectItem +from PySide6.QtWidgets import QWidget, QVBoxLayout, QSizePolicy, QGraphicsView +from PySide6.QtWidgets import QGraphicsScene, QGraphicsEllipseItem, QGraphicsRectItem, QGraphicsLineItem from PySide6.QtCore import Qt, QPointF, QRectF, QPointF -from PySide6.QtGui import QPixmap, QBrush, QColor, QImage +from PySide6.QtGui import QPixmap, QBrush, QColor, QImage, QPen -skeleton_grid = [(0, 1), (1, 2), (2, 3), (2, 4), (2, 5)] class Skeleton(QGraphicsRectItem): + skeleton_grid = [(0, 1), (1, 2), (1, 3), (1, 4), (2, 5)] def __init__(self, x, y, width, height, keypoint_coordinates, brush): super().__init__(x, y, width, height) - - for kc in keypoint_coordinates: - kx = kc[0] - x - ky = kc[1] - y + b = QBrush(QColor(255,255,255,64)) + p = QPen(brush.color()) + markersize = 5 + p.setWidth(1) + p.setStyle(Qt.PenStyle.SolidLine) + self.setBrush(b) + for i, kc in enumerate(keypoint_coordinates): + kx = kc[0] - x - markersize/2 + ky = kc[1] - y - markersize/2 kp = QGraphicsEllipseItem(kx, ky, 5, 5, self) kp.setBrush(brush) - # for sg in skeleton_grid: - # gsx = keypoint_coordinates[sg[0], ] + for i, sg in enumerate(self.skeleton_grid): + gsx = keypoint_coordinates[sg[0], 0] - x # grid start x + gsy = keypoint_coordinates[sg[0], 1] - y + gex = keypoint_coordinates[sg[1], 0] - x + gey = keypoint_coordinates[sg[1], 1] - y + l = QGraphicsLineItem(gsx, gsy, gex, gey, self) + p.color().setAlpha(255/(i+1)) + l.setPen(p) # self.setBrush(brush) # self.setAcceptHoverEvents(True) # Enable hover events if needed self.setFlags(QGraphicsRectItem.ItemIsSelectable) @@ -41,7 +53,7 @@ class SkeletonWidget(QWidget): self._view = QGraphicsView() self._view.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) self._view.setMouseTracking(True) - + self._view.setScene(self._scene) lyt = QVBoxLayout() lyt.addWidget(self._view) self.setLayout(lyt) @@ -54,21 +66,14 @@ class SkeletonWidget(QWidget): item = Skeleton(boxx, boxy, boxw, boxh, coords, brush) item.setData(0, id) self._scene.addItem(item) - self._scene.addRect(0,0, 10, 10) def addSkeletons(self, coordinates:np.array, detection_ids:np.array, brush:QBrush): - image_rect = self._pixmapitem.boundingRect() if self._pixmapitem is not None else QRectF(0,0,0,0) num_detections = coordinates.shape[0] for i in range(num_detections): - x = coordinates[i, 0] - y = coordinates[i, 1] c = brush.color() c.setAlpha(int(i * 255 / num_detections)) brush.setColor(c) - item = Skeleton(brush=brush) - item.setData(0, track_ids[i]) - item.setData(1, detection_ids[i]) - item = self._scene.addItem(item) + self.addSkeleton(coordinates[i,:,:], detection_ids[i], brush=brush) logging.debug("DetectionView: Number of items in scene: %i", len(self._scene.items())) @@ -103,8 +108,11 @@ def main(): view = SkeletonWidget() # view.signals.itemsSelected.connect(items_selected) layout.addWidget(view) - view.addSkeleton(focus_coords[10,:,:], focus_ids[10], focus_brush) - print(view._scene.items()) + # view.addSkeleton(focus_coords[10,:,:], focus_ids[10], focus_brush) + count = 20 + view.addSkeletons(focus_coords[:count,:,:], focus_ids[:count], focus_brush) + view.addSkeletons(scnd_coords[:count,:,:], scnd_ids[:count], second_brush) + # view.addSkeletons(focus_coords[:10,:,:], focus_ids[:10], focus_brush) # view.addDetections(focus_coords, focus_tracks, focus_ids, focus_brush) # view.addDetections(scnd_coords, scnd_tracks, scnd_ids, second_brush) diff --git a/fixtracks/widgets/tracks.py b/fixtracks/widgets/tracks.py index eb69993..4d797b9 100644 --- a/fixtracks/widgets/tracks.py +++ b/fixtracks/widgets/tracks.py @@ -13,6 +13,7 @@ from fixtracks.utils.reader import PickleLoader from fixtracks.utils.writer import PickleWriter from fixtracks.widgets.detectionview import DetectionView from fixtracks.widgets.detectiontimeline import DetectionTimeline +from fixtracks.widgets.skeleton import SkeletonWidget class PoseTableModel(QAbstractTableModel): column_header = ["frame", "track"] @@ -352,6 +353,9 @@ class FixTracks(QWidget): self._controls_widget.fwd.connect(self.on_forward) self._controls_widget.back.connect(self.on_backward) + self._skeleton = SkeletonWidget() + self._skeleton.setMaximumSize(QSize(400, 400)) + self._saveBtn = QPushButton("Save") self._saveBtn.setShortcut("Ctrl+S") self._saveBtn.setEnabled(False) @@ -381,10 +385,14 @@ 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) vbox = QVBoxLayout() vbox.addLayout(timelinebox) - vbox.addWidget(self._controls_widget, stretch=1, alignment=Qt.AlignmentFlag.AlignCenter) + vbox.addLayout(cntrlBox) vbox.addLayout(btnBox) container = QWidget() container.setLayout(vbox)