diff --git a/fixtracks/widgets/skeleton.py b/fixtracks/widgets/skeleton.py new file mode 100644 index 0000000..4e8038e --- /dev/null +++ b/fixtracks/widgets/skeleton.py @@ -0,0 +1,116 @@ +import logging +import numpy as np + +from PySide6.QtWidgets import QWidget, QVBoxLayout, QSizePolicy, QGraphicsView, QGraphicsItem, QGraphicsScene, QGraphicsEllipseItem, QGraphicsRectItem +from PySide6.QtCore import Qt, QPointF, QRectF, QPointF +from PySide6.QtGui import QPixmap, QBrush, QColor, QImage + +skeleton_grid = [(0, 1), (1, 2), (2, 3), (2, 4), (2, 5)] + +class Skeleton(QGraphicsRectItem): + + 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 + kp = QGraphicsEllipseItem(kx, ky, 5, 5, self) + kp.setBrush(brush) + # for sg in skeleton_grid: + # gsx = keypoint_coordinates[sg[0], ] + # self.setBrush(brush) + # self.setAcceptHoverEvents(True) # Enable hover events if needed + self.setFlags(QGraphicsRectItem.ItemIsSelectable) + + # def mousePressEvent(self, event): + # self.signals.clicked.emit(self.data(0), QPointF(event.scenePos().x(), event.scenePos().y())) + + # def hoverEnterEvent(self, event): + # self.signals.hover.emit(self.data(0), QPointF(event.scenePos().x(), event.scenePos().y())) + # super().hoverEnterEvent(event) + + +class SkeletonWidget(QWidget): + + def __init__(self, parent=None): + super().__init__(parent) + self._img = None + self._scene = QGraphicsScene() + # self.setRenderHint(QGraphicsView.RenderFlag.Ren Antialiasing) + self._view = QGraphicsView() + self._view.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding) + self._view.setMouseTracking(True) + + lyt = QVBoxLayout() + lyt.addWidget(self._view) + self.setLayout(lyt) + + def addSkeleton(self, coords, id, brush): + boxx = np.min(coords[:,0]) + boxy = np.min(coords[:,1]) + boxw = np.max(coords[:, 0]) - boxx + boxh = np.max(coords[:, 1]) - boxy + 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) + logging.debug("DetectionView: Number of items in scene: %i", len(self._scene.items())) + + +def main(): + import pickle + import numpy as np + from IPython import embed + from PySide6.QtWidgets import QApplication + from fixtracks.info import PACKAGE_ROOT + + datafile = PACKAGE_ROOT / "data/merged_small.pkl" + print(datafile) + with open(datafile, "rb") as f: + df = pickle.load(f) + + focus_brush = QBrush(QColor.fromString("red")) + second_brush = QBrush(QColor.fromString("blue")) + + scnd_coords = np.stack(df.keypoints[(df.track == 2)].values,).astype(np.float32)[:,:,:] + scnd_tracks = df.track[df.track == 2].values + scnd_ids = df.track[(df.track == 2)].index.values + + focus_coords = np.stack(df.keypoints[df.track == 1].values,).astype(np.float32)[:,:,:] + focus_tracks = df.track[df.track == 1].values + focus_ids = df.track[(df.track == 2)].index.values + + app = QApplication([]) + window = QWidget() + window.setMinimumSize(200, 200) + layout = QVBoxLayout() + + 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.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) + window.setLayout(layout) + window.show() + app.exec() + +if __name__ == "__main__": + main() \ No newline at end of file