diff --git a/fixtracks/widgets/skeleton.py b/fixtracks/widgets/skeleton.py index e2fd2bb..6e7e80a 100644 --- a/fixtracks/widgets/skeleton.py +++ b/fixtracks/widgets/skeleton.py @@ -10,9 +10,11 @@ from fixtracks.widgets.detectionview 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 @@ -37,6 +39,17 @@ class Skeleton(QGraphicsRectItem): # self.setAcceptHoverEvents(True) # Enable hover events if needed self.setFlags(QGraphicsRectItem.ItemIsSelectable) + @property + def length(self): + dist = 0.0 + for i in range(1, len(self.bodyaxis)): + si = self.bodyaxis[i-1] + ei = self.bodyaxis[i] + sk = self._keypoints[si, :] + ek = self._keypoints[ei, :] + dist += np.sqrt((ek[0] - sk[0])**2 + (ek[1] - sk[1])**2) + return dist + # def mousePressEvent(self, event): # self.signals.clicked.emit(self.data(0), QPointF(event.scenePos().x(), event.scenePos().y())) @@ -68,7 +81,7 @@ class SkeletonWidget(QWidget): font.setPointSize(9) self._info_label = QLabel("") self._info_label.setFont(font) - + lyt = QVBoxLayout() lyt.addWidget(self._view) lyt.addWidget(self._info_label) @@ -82,7 +95,11 @@ class SkeletonWidget(QWidget): def updateInfo(self, index): if index > -1: s = self._skeletons[index] - self._info_label.setText(f"Detection id {s.data(DetectionData.ID.value)}, track {s.data(DetectionData.TRACK_ID.value)} on frame {s.data(DetectionData.FRAME.value)}") + l = s.length + i = s.data(DetectionData.ID.value) + t = s.data(DetectionData.TRACK_ID.value) + f = s.data(DetectionData.FRAME.value) + self._info_label.setText(f"Id {i}, track {t} on frame {f}, length {l:.1f} px") else: self._info_label.setText("") @@ -119,6 +136,8 @@ class SkeletonWidget(QWidget): def addSkeleton(self, coords, detection_id, frame, track, 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 @@ -186,14 +205,9 @@ def main(): 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_frames = df.track[df.track == 1].values focus_ids = df.track[(df.track == 2)].index.values app = QApplication([]) @@ -209,7 +223,8 @@ def main(): 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(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)