239 lines
9.4 KiB
Python
239 lines
9.4 KiB
Python
import logging
|
|
import numpy as np
|
|
|
|
from PySide6.QtWidgets import QWidget, QVBoxLayout, QSizePolicy, QGraphicsView, QSlider, QPushButton, QLabel
|
|
from PySide6.QtWidgets import QGraphicsScene, QGraphicsEllipseItem, QGraphicsRectItem, QGraphicsLineItem
|
|
from PySide6.QtCore import Qt
|
|
from PySide6.QtGui import QBrush, QColor, QPen, QPainter, QFont
|
|
|
|
from fixtracks.utils.enums import DetectionData
|
|
|
|
class Skeleton(QGraphicsRectItem):
|
|
skeleton_grid = [(0, 1), (1, 2), (1, 3), (1, 4), (2, 5)]
|
|
bodyaxis = [0, 1, 2, 5]
|
|
|
|
def __init__(self, x, y, width, height, keypoint_coordinates, brush):
|
|
super().__init__(x, y, width, height)
|
|
self._keypoints = keypoint_coordinates
|
|
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)
|
|
|
|
@property
|
|
def length(self):
|
|
bodykpts = self._keypoints[self.bodyaxis, :]
|
|
dist = np.sum(np.sqrt(np.sum(np.diff(bodykpts, axis=0)**2, axis=1)), axis=0)
|
|
return dist
|
|
|
|
# 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._maxx = 0.
|
|
self._minx = 0.
|
|
self._maxy = 0.
|
|
self._miny = 0.
|
|
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)
|
|
|
|
font = QFont()
|
|
font.setPointSize(9)
|
|
self._info_label = QLabel("")
|
|
self._info_label.setFont(font)
|
|
|
|
lyt = QVBoxLayout()
|
|
lyt.addWidget(self._view)
|
|
lyt.addWidget(self._info_label)
|
|
lyt.addWidget(self._slider)
|
|
self.setLayout(lyt)
|
|
|
|
self._view.fitInView(self._scene.sceneRect(), Qt.AspectRatioMode.KeepAspectRatio)
|
|
self._scene.changed.connect(lambda: self._view.fitInView(self._scene.sceneRect(),
|
|
Qt.AspectRatioMode.KeepAspectRatio))
|
|
|
|
def updateInfo(self, index):
|
|
if index > -1:
|
|
s = self._skeletons[index]
|
|
l = s.length
|
|
i = s.data(DetectionData.ID.value)
|
|
t = s.data(DetectionData.TRACK_ID.value)
|
|
f = s.data(DetectionData.FRAME.value)
|
|
sc = s.data(DetectionData.SCORE.value)
|
|
self._info_label.setText(f"Id {i}, track {t} on frame {f}, length {l:.1f} px, confidence {sc:.2f}")
|
|
else:
|
|
self._info_label.setText("")
|
|
|
|
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()
|
|
self.updateInfo(-1)
|
|
|
|
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)
|
|
self.updateInfo(self._slider.value())
|
|
|
|
def update(self):
|
|
logging.debug("SkeletonWidget: update")
|
|
if len(self._skeletons) > 0:
|
|
self._scene.addItem(self._skeletons[-1])
|
|
self._current_skeleton = self._skeletons[-1]
|
|
self.updateInfo(len(self._skeletons)-1)
|
|
self._slider.setMaximum(len(self._skeletons))
|
|
self._slider.setMinimum(0)
|
|
self._slider.setValue(len(self._skeletons))
|
|
self._slider.setEnabled(len(self._skeletons) > 0)
|
|
self._scene.setSceneRect(self._minx, self._miny, self._maxx - self._minx, self._maxy - self._miny)
|
|
self._view.fitInView(self._scene.sceneRect(), Qt.AspectRatioMode.KeepAspectRatio)
|
|
|
|
def addSkeleton(self, coords, detection_id, frame, track, score, brush, update=True):
|
|
def check_extent(x, y, w, h):
|
|
if x == 0 and y == 0:
|
|
return
|
|
if len(self._skeletons) == 0:
|
|
self._minx = x
|
|
self._maxx = x + w
|
|
self._miny = y
|
|
self._maxy = y + h
|
|
else:
|
|
if x < self._minx:
|
|
self._minx = x
|
|
if y > self._maxy:
|
|
self._maxy = y
|
|
if x + w > self._maxx:
|
|
self._maxx = x + w
|
|
if y + h > self._maxy:
|
|
self._maxy = y + h
|
|
|
|
boxx = np.min(coords[:,0])
|
|
boxy = np.min(coords[:,1])
|
|
boxw = np.max(coords[:, 0]) - boxx
|
|
boxh = np.max(coords[:, 1]) - boxy
|
|
check_extent(boxx, boxy, boxw, boxh)
|
|
item = Skeleton(boxx, boxy, boxw, boxh, coords, brush)
|
|
item.setData(DetectionData.ID.value, detection_id)
|
|
item.setData(DetectionData.TRACK_ID.value, track)
|
|
item.setData(DetectionData.FRAME.value, frame)
|
|
item.setData(DetectionData.SCORE.value, score)
|
|
self._skeletons.append(item)
|
|
if update:
|
|
self.update()
|
|
|
|
def addSkeletons(self, coordinates:np.ndarray, detection_ids:np.ndarray,
|
|
frames:np.ndarray, tracks:np.ndarray, scores:np.ndarray,
|
|
brush:QBrush):
|
|
num_detections = 0 if coordinates is None else coordinates.shape[0]
|
|
logging.debug("SkeletonWidget: add %i Skeletons", num_detections)
|
|
if num_detections < 1:
|
|
return
|
|
sorting = np.argsort(frames)
|
|
coordinates = coordinates[sorting,:, :]
|
|
detection_ids = detection_ids[sorting]
|
|
frames = frames[sorting]
|
|
tracks = tracks[sorting]
|
|
scores = scores[sorting]
|
|
for i in range(num_detections):
|
|
self.addSkeleton(coordinates[i,:,:], detection_ids[i], frames[i],
|
|
tracks[i], scores[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"))
|
|
focus_coords = np.stack(df.keypoints[df.track == 1].values,).astype(np.float32)[:,:,:]
|
|
focus_tracks = df.track[df.track == 1].values
|
|
focus_frames = 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_frames[:count], focus_tracks[: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() |