[detectionview] widget to display detections on top of image

fixtracks/widgets/detectionview.py

@@ -0,0 +1,120 @@
+import logging
+import numpy as np
+import pandas as pd 
+
+from PySide6.QtWidgets import QWidget, QVBoxLayout, QSizePolicy, QGraphicsView, QGraphicsScene, QGraphicsEllipseItem
+from PySide6.QtCore import Qt, QPointF
+from PySide6.QtGui import QPixmap, QBrush, QColor, QImage
+
+
+from fixtracks.info import PACKAGE_ROOT
+from fixtracks.utils.reader  import PickleLoader
+from fixtracks.utils.signals import DetectionSignals
+
+class Detection(QGraphicsEllipseItem):
+    signals = DetectionSignals()
+
+    def __init__(self, x, y, width, height, brush):
+        super().__init__(x, y, width, height)
+        self.setBrush(brush)
+        self.setAcceptHoverEvents(True)  # Enable hover events if needed
+
+    def mousePressEvent(self, event):
+        self.signals.clicked.emit(self.data(0), QPointF(event.scenePos().x(), event.scenePos().y()))
+        # print(f"Rectangle clicked at: {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 DetectionView(QWidget):
+
+    def __init__(self, parent=None):
+        super().__init__(parent)
+        self._img = None
+        self._pixmapitem = None
+        self._scene = None
+        self._view = QGraphicsView()
+        self._view.setSizePolicy(QSizePolicy.Policy.Expanding, QSizePolicy.Policy.Expanding)
+        self._view.setMouseTracking(True)
+        self._items = []
+
+        lyt = QVBoxLayout()
+        lyt.addWidget(self._view)
+        self.setLayout(lyt)
+
+    def setImage(self, image: QImage):
+        self._img = image
+        self._scene = QGraphicsScene()
+        self._pixmapitem = self._scene.addPixmap(QPixmap.fromImage(self._img))
+        self._view.setScene(self._scene)
+        self._view.fitInView(self._scene.sceneRect(), aspectRadioMode=Qt.AspectRatioMode.KeepAspectRatio)
+        # self._view.show()
+
+    def clearDetections(self):
+        for item in self._items:
+            self._scene.removeItem(item)
+
+    def addDetections(self, coordinates:np.array, ids:np.array, brush:QBrush):
+        image_rect = self._pixmapitem.boundingRect()
+        for i in range(coordinates.shape[0]):
+            x = coordinates[i, 0]
+            y = coordinates[i, 1]
+            item = Detection(image_rect.left() + x, image_rect.top() + y, 10, 10, brush=brush)
+            item.setData(0, ids[i])
+            item = self._scene.addItem(item)
+            self._items.append(item)
+
+    def fit_image_to_view(self):
+        """Scale the image to fit the QGraphicsView."""
+        if self._pixmapitem is not None:
+            self._view.fitInView(self._pixmapitem, Qt.KeepAspectRatio)
+
+    def resizeEvent(self, event):
+        """Handle window resizing to fit the image."""
+        super().resizeEvent(event)
+        self.fit_image_to_view()
+
+def main():
+    import pickle
+    import numpy as np
+    from IPython import embed
+    from PySide6.QtWidgets import QApplication
+
+    datafile = PACKAGE_ROOT / "data/merged_small.pkl"
+    imgfile = PACKAGE_ROOT / "data/merged.png"
+    print(datafile)
+    with open(datafile, "rb") as f:
+        df = pickle.load(f)
+    img = QImage(imgfile)
+    focus_brush = QBrush(QColor.fromString("red"))
+    second_brush = QBrush(QColor.fromString("blue"))
+    background_brush = QBrush(QColor.fromString("white"))
+
+    bg_coords = np.stack(df.keypoints[(df.track != 1) & (df.track != 2)].values,).astype(np.float32)[:,0,:]
+    bg_ids = df.track[(df.track != 1) & (df.track != 2)].values
+
+    scnd_coords = np.stack(df.keypoints[(df.track == 2)].values,).astype(np.float32)[:,0,:]
+    scnd_ids = df.track[df.track == 2].values
+
+    focus_coords = np.stack(df.keypoints[df.track == 1].values,).astype(np.float32)[:,0,:]
+    focus_ids = df.track[df.track == 1].values
+
+    app = QApplication([])
+    window = QWidget()
+    window.setMinimumSize(200, 200)
+    layout = QVBoxLayout()
+
+    view = DetectionView()
+    layout.addWidget(view)
+    view.setImage(img)
+    view.addDetections(bg_coords, bg_ids, background_brush)
+    view.addDetections(focus_coords, focus_ids, focus_brush)
+    view.addDetections(scnd_coords, scnd_ids, second_brush)
+    window.setLayout(layout)
+    # window.show()
+    app.exec()
+
+if __name__ == "__main__":
+    main()