fixtracks/fixtracks/widgets/skeleton.py

import logging
import numpy as np

from PySide6.QtWidgets import QWidget, QVBoxLayout, QSizePolicy, QGraphicsView, QSlider, QPushButton
from PySide6.QtWidgets import QGraphicsScene, QGraphicsEllipseItem, QGraphicsRectItem, QGraphicsLineItem
from PySide6.QtCore import Qt
from PySide6.QtGui import QBrush, QColor, QPen, QPainter


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)
        skeleton_pen = QPen(brush.color())
        skeleton_pen.setWidthF(1.0)
        skeleton_marker = 5
        bg_brush = QBrush(QColor(127,127,127,32))
        bg_pen = QPen(QColor(127, 127, 127, 255))
        bg_pen.setWidth(0.5)
        bg_pen.setStyle(Qt.PenStyle.DashDotLine)
        self.setBrush(bg_brush)
        self.setPen(bg_pen)
        for i, kc in enumerate(keypoint_coordinates):
            kx = kc[0] - skeleton_marker/2
            ky = kc[1] - skeleton_marker/2
            kp = QGraphicsEllipseItem(kx, ky, 5, 5, self)
            kp.setBrush(brush)
        for i, sg in enumerate(self.skeleton_grid):
            gsx = keypoint_coordinates[sg[0], 0]  # grid start x
            gsy = keypoint_coordinates[sg[0], 1]
            gex = keypoint_coordinates[sg[1], 0]
            gey = keypoint_coordinates[sg[1], 1]
            QGraphicsLineItem(gsx, gsy, gex, gey, self)

        # 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):
    #FIXME May be more efficient to set hide and show skeletons instead of adding and removing them when moving the slider
    def __init__(self, parent=None):
        super().__init__(parent)
        self._skeletons = []
        self._current_skeleton = None
        self._slider = QSlider(Qt.Orientation.Horizontal)
        self._slider.sliderMoved.connect(self.on_sliderMoved)
        self._scene = QGraphicsScene()
        self._view = QGraphicsView()
        self._view.setRenderHint(QPainter.Antialiasing)
        self._view.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
        self._view.setMouseTracking(True)
        self._view.setScene(self._scene)

        lyt = QVBoxLayout()
        lyt.addWidget(self._view)
        lyt.addWidget(self._slider)
        self.setLayout(lyt)

    def clear(self):
        for i in range(len(self._skeletons)):
            item = self._skeletons.pop()
            if item.scene() == self._scene:
                self._scene.removeItem(item)
            del item
        self._skeletons = []
        self.update()

    def on_sliderMoved(self):
        if self._current_skeleton is not None and self._current_skeleton.scene() == self._scene:
            self._scene.removeItem(self._current_skeleton)
        if self._slider.value() < len(self._skeletons):
            self._current_skeleton = self._skeletons[self._slider.value()]
            self._scene.addItem(self._current_skeleton)

    def update(self):
        if len(self._skeletons) > 0:
            self._scene.addItem(self._skeletons[-1])
            self._current_skeleton = self._skeletons[-1]
        self._slider.setMaximum(len(self._skeletons))
        self._slider.setMinimum(0)
        self._slider.setValue(len(self._skeletons))
        if len(self._skeletons) == 0:
            self._slider.setEnabled(False)

    def addSkeleton(self, coords, detection_id, brush, update=True):
        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, detection_id)
        self._skeletons.append(item)
        if update:
            self.update()

    def addSkeletons(self, coordinates, detection_ids, brush):
        num_detections = coordinates.shape[0]
        for i in range(num_detections):
            self.addSkeleton(coordinates[i,:,:], detection_ids[i], brush=brush, update=False)
        self.update()

    # def addSkeleton(self, coords, detection_id, brush):
    #     item = self.createSkeleton(coords, detection_id, brush)
    #     self._skeletons.append(item)
    #     self._scene.addItem(item)

    # def addSkeletons(self, coordinates:np.array, detection_ids:np.array, brush:QBrush):
    #     skeletons = self.createSkeletons(coordinates, detection_ids, brush)
    #     self._skeletons.extend(skeletons)
    #     for s in skeletons:
    #         self._scene.addItem(s)


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)
    btn = QPushButton("clear")
    btn.clicked.connect(view.clear)
    layout.addWidget(btn)
    # view.addSkeleton(focus_coords[10,:,:], focus_ids[10], focus_brush)
    count = 100
    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)
    window.setLayout(layout)
    window.show()
    app.exec()

if __name__ == "__main__":
    main()