From 4a766557663a4e212b3175912d74535056ca1d89 Mon Sep 17 00:00:00 2001 From: Jan Grewe Date: Thu, 27 Feb 2025 19:59:46 +0100 Subject: [PATCH] [detections] add score to items --- fixtracks/utils/enums.py | 1 + fixtracks/widgets/detectionview.py | 11 ++++++++--- fixtracks/widgets/skeleton.py | 12 ++++++++---- fixtracks/widgets/tracks.py | 15 ++++++++++++--- 4 files changed, 29 insertions(+), 10 deletions(-) diff --git a/fixtracks/utils/enums.py b/fixtracks/utils/enums.py index 3e378cf..0bdd7bd 100644 --- a/fixtracks/utils/enums.py +++ b/fixtracks/utils/enums.py @@ -8,6 +8,7 @@ class DetectionData(Enum): COORDINATES = 2 TRACK_ID = 3 USERLABELED = 4 + SCORE = 5 class Tracks(Enum): TRACKONE = 1 diff --git a/fixtracks/widgets/detectionview.py b/fixtracks/widgets/detectionview.py index feb1a00..059be63 100644 --- a/fixtracks/widgets/detectionview.py +++ b/fixtracks/widgets/detectionview.py @@ -10,6 +10,7 @@ from fixtracks.utils.signals import DetectionSignals, DetectionViewSignals, Dete from fixtracks.utils.enums import DetectionData, Tracks from fixtracks.utils.trackingdata import TrackingData + class Detection(QGraphicsEllipseItem): signals = DetectionSignals() @@ -138,10 +139,11 @@ class DetectionView(QWidget): coordinates = self._data.coordinates(selection=True) centercoordinates = self._data.centerOfGravity(selection=True) userlabeled = self._data.selectedData("userlabeled") + scores = self._data.selectedData("confidence") image_rect = self._pixmapitem.boundingRect() if self._pixmapitem is not None else QRectF(0,0,0,0) - for i, (id, f, t, l) in enumerate(zip(ids, frames, tracks, userlabeled)): + for i, (id, f, t, l, s) in enumerate(zip(ids, frames, tracks, userlabeled, scores)): c = Tracks.fromValue(t).toColor() if keypoint >= 0: x = coordinates[i, keypoint, 0] @@ -155,7 +157,10 @@ class DetectionView(QWidget): item.setData(DetectionData.ID.value, id) item.setData(DetectionData.COORDINATES.value, coordinates[i, :, :]) item.setData(DetectionData.FRAME.value, f) - item.setData(DetectionData.USERLABELED.value, l) + # item.setData(DetectionData.USERLABELED.value, l) + item.setData(DetectionData.SCORE.value, s) + print(s) + print(item.data(DetectionData.SCORE.value)) item = self._scene.addItem(item) def fit_image_to_view(self): @@ -214,7 +219,7 @@ def main(): view.setImage(img) view.addDetections(bg_coords, bg_tracks, bg_ids, background_brush) view.addDetections(focus_coords, focus_tracks, focus_ids, focus_brush) - view.addDetections(scnd_coords, scnd_tracks, scnd_ids, second_brush) + view.addDetections(scnd_coords, scnd_tracks, scnd_ids, second_brush) window.setLayout(layout) window.show() app.exec() diff --git a/fixtracks/widgets/skeleton.py b/fixtracks/widgets/skeleton.py index 59a9822..290a5c2 100644 --- a/fixtracks/widgets/skeleton.py +++ b/fixtracks/widgets/skeleton.py @@ -94,7 +94,8 @@ class SkeletonWidget(QWidget): 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") + 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("") @@ -129,7 +130,7 @@ class SkeletonWidget(QWidget): 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, brush, update=True): + 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 @@ -157,12 +158,14 @@ class SkeletonWidget(QWidget): 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, brush:QBrush): + 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: @@ -172,9 +175,10 @@ class SkeletonWidget(QWidget): 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], brush=brush, update=False) + tracks[i], scores[i], brush=brush, update=False) self.update() # def addSkeleton(self, coords, detection_id, brush): diff --git a/fixtracks/widgets/tracks.py b/fixtracks/widgets/tracks.py index 0da68a1..40badb2 100644 --- a/fixtracks/widgets/tracks.py +++ b/fixtracks/widgets/tracks.py @@ -65,8 +65,10 @@ class FixTracks(QWidget): combo_layout.addWidget(QLabel("Window width:")) combo_layout.addWidget(self._windowspinner) combo_layout.addWidget(QLabel("frames")) + combo_layout.addItem(QSpacerItem(10, 10, QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed)) combo_layout.addWidget(QLabel("Keypoint:")) combo_layout.addWidget(self._keypointcombo) + combo_layout.addItem(QSpacerItem(10, 10, QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Fixed)) combo_layout.addWidget(QLabel("Jump to frame:")) combo_layout.addWidget(self._gotoframe) combo_layout.addWidget(self._gotobtn) @@ -168,15 +170,19 @@ class FixTracks(QWidget): self._detectionView.setImage(img) def update(self): + kp = self._keypointcombo.currentText().lower() + if len(kp) == 0: + return + kpi = -1 if "center" in kp else int(kp) + start_frame = self._currentWindowPos stop_frame = start_frame + self._currentWindowWidth + self._timeline.setWindow(start_frame / self._maxframes, self._currentWindowWidth/self._maxframes) logging.debug("Tracks:update: Updating View for detection range %i, %i frames", start_frame, stop_frame) self._data.setSelectionRange("frame", start_frame, stop_frame) self._controls_widget.setWindow(start_frame, stop_frame) - kp = self._keypointcombo.currentText().lower() - kpi = -1 if "center" in kp else int(kp) self._detectionView.updateDetections(kpi) @property @@ -358,6 +364,7 @@ class FixTracks(QWidget): tracks = np.zeros(len(detections), dtype=int) ids = np.zeros_like(tracks) frames = np.zeros_like(tracks) + scores = np.zeros_like(tracks) coordinates = None if len(detections) > 0: c = detections[0].data(DetectionData.COORDINATES.value) @@ -368,10 +375,12 @@ class FixTracks(QWidget): ids[i] = d.data(DetectionData.ID.value) frames[i] = d.data(DetectionData.FRAME.value) coordinates[i, :, :] = d.data(DetectionData.COORDINATES.value) + scores[i] = d.data(DetectionData.SCORE.value) + print(scores[i]) self._data.setSelection(ids) self._controls_widget.setSelectedTracks(tracks) self._skeleton.clear() - self._skeleton.addSkeletons(coordinates, ids, frames, tracks, QBrush(QColor(10, 255, 65, 255))) + self._skeleton.addSkeletons(coordinates, ids, frames, tracks, scores, QBrush(QColor(10, 255, 65, 255))) def moveWindow(self, stepsize): logging.info("Tracks.moveWindow: move window with stepsize %.2f", stepsize)