diff --git a/fixtracks/widgets/detectionview.py b/fixtracks/widgets/detectionview.py new file mode 100644 index 0000000..5800cbd --- /dev/null +++ b/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()