Changes for seamless combatibility with pyidi analyses returns

Author: klemengit

pyidi/GUIs/__init__.py

@@ -1,11 +1,13 @@
+import typing
+
 try:
     import PyQt6
 
     HAS_PYQT6 = True
 except ImportError:
     HAS_PYQT6 = False
 
-if HAS_PYQT6:
+if HAS_PYQT6 or typing.TYPE_CHECKING:
     from .subset_selection import SelectionGUI
     from .result_viewer import ResultViewer
 else:

pyidi/GUIs/result_viewer.py

@@ -7,11 +7,47 @@
 import sys
 
 class ResultViewer(QtWidgets.QMainWindow):
-    def __init__(self, video, displacements, grid, fps=30, magnification=1, point_size=10, colormap="cool"):
+    def __init__(self, video, displacements, points, fps=30, magnification=1, point_size=10, colormap="cool"):
+        """
+        The results from the pyidi analysis can directly be passed to this class:
+        
+        - ``video``: can be a ``VideoReader`` object (or numpy array of correct shape).
+        - ``displacements``: directly the return from the ``get_displacements`` method.
+        - ``points``: the points used for the analysis, which were passed to the ``set_points`` method.
+
+        Parameters
+        ----------
+        video : np.ndarray or VideoReader
+            Array of shape (n_frames, height, width) containing the video frames.
+        displacements : np.ndarray
+            Array of shape (n_frames, n_points, 2) containing the displacement vectors.
+        points : np.ndarray
+            Array of shape (n_points, 2) containing the grid points.
+        fps : int, optional
+            Frames per second for the video playback, by default 30.
+        magnification : int, optional
+            Magnification factor for the displacements, by default 1.
+        point_size : int, optional
+            Size of the points in pixels, by default 10.
+        colormap : str, optional
+            Name of the colormap to use for the arrows, by default "cool".
+        """
+        # Create QApplication if it doesn't exist
+        app = QtWidgets.QApplication.instance()
+        if app is None:
+            app = QtWidgets.QApplication([])
+        
         super().__init__()
-        self.video = video
-        self.displacements = displacements
-        self.grid = grid
+        
+        # Coordinate transformation to match viewer function behavior
+        from ..video_reader import VideoReader
+        if isinstance(video, VideoReader):
+            self.video = video.get_frames()
+        else:
+            self.video = video
+
+        self.displacements = displacements[:, :, ::-1]  # Flip x,y coordinates
+        self.grid = points[:, ::-1] + 0.5  # Flip x,y coordinates
         self.fps = fps
         self.magnification = magnification
         self.points_size = point_size
@@ -25,6 +61,14 @@ def __init__(self, video, displacements, grid, fps=30, magnification=1, point_si
 
         self.init_ui()
         self.update_frame()
+        
+        # Start the GUI
+        self.show()
+        # Only call sys.exit if not in IPython
+        if not hasattr(sys, 'ps1'):  # Not interactive
+            sys.exit(app.exec())
+        else:
+            app.exec()  # Don't raise SystemExit in IPython
 
     def init_ui(self):
         # Style
@@ -173,6 +217,7 @@ def init_ui(self):
         # === Finalize ===
         self.setCentralWidget(central_widget)
         self.setWindowTitle("Displacement Viewer")
+        self.resize(800, 600)
 
 
     def toggle_playback(self):
@@ -337,22 +382,10 @@ def viewer(frames, displacements, points, fps=30, magnification=1, point_size=10
     colormap : str, optional
         Name of the colormap to use for the arrows, by default "cool".
     """
-    points = points[:, ::-1]
-    displacements = displacements[:, :, ::-1]
-
-    app = QtWidgets.QApplication.instance()
-    if app is None:
-        app = QtWidgets.QApplication([])
-    
-    win = ResultViewer(frames, displacements, points, fps=fps, magnification=magnification, point_size=point_size, colormap=colormap)
-    win.resize(800, 600)
-    win.show()
-    
-    # Only call sys.exit if not in IPython
-    if not hasattr(sys, 'ps1'):  # Not interactive
-        sys.exit(app.exec())
-    else:
-        app.exec()  # Don't raise SystemExit in IPythonys
+    # This function is now just a wrapper for backward compatibility
+    # The ResultViewer class handles everything internally
+    ResultViewer(frames, displacements, points, fps=fps, magnification=magnification, 
+                point_size=point_size, colormap=colormap)
 
 
 if __name__ == "__main__":
@@ -362,5 +395,6 @@ def viewer(frames, displacements, points, fps=30, magnification=1, point_size=10
     displacements = 2 * (np.random.rand(n_points, n_frames, 2) - 0.5)
     grid = np.stack(np.meshgrid(np.linspace(50, 350, int(np.sqrt(n_points))),
                                 np.linspace(50, 250, int(np.sqrt(n_points)))), axis=-1).reshape(-1, 2)[:n_points]
-    grid = grid[:, ::-1]
-    viewer(frames, displacements, grid)
+    
+    # Now you can directly call ResultViewer (no need to flip coordinates here since it's done internally)
+    ResultViewer(frames, displacements, grid)

pyidi/GUIs/subset_selection.py

@@ -41,6 +41,15 @@ def mouseDragEvent(self, ev, axis=None):
 
 class SelectionGUI(QtWidgets.QMainWindow):
     def __init__(self, video):
+        """Initialize the selection GUI for manual subset selection.
+
+        To extract the points, use the ``get_points`` method or the ``points`` attribute.
+        
+        Parameters
+        ----------
+        video : VideoReader or np.ndarray
+            The video to be analyzed. If a VideoReader object, it should be initialized with the video file.
+        """
         app = QtWidgets.QApplication.instance()
         if app is None:
             app = QtWidgets.QApplication([])
@@ -156,7 +165,11 @@ def __init__(self, video):
         self.pg_widget.scene().sigMouseClicked.connect(self.on_mouse_click)
 
         # Set the initial image
-        self.image_item.setImage(video)
+        from ..video_reader import VideoReader
+        if isinstance(video, VideoReader):
+            self.frame = video.get_frame(0)
+            
+        self.image_item.setImage(self.frame.T) # axis 0 is x, while image axis 0 is y
 
         # Ensure method-specific widgets are visible on startup
         self.method_selected(self.button_group.checkedId())
@@ -204,7 +217,7 @@ def create_help_button(self, tooltip_text: str) -> QtWidgets.QToolButton:
     def ui_graphics(self):
         # Image viewer
         self.pg_widget = GraphicsLayoutWidget()
-        self.view = BrushViewBox(parent_gui=self, lockAspect=True)
+        self.view = BrushViewBox(parent_gui=self, lockAspect=True, invertY=True)
         self.pg_widget.addItem(self.view)
 
 
@@ -763,7 +776,12 @@ def set_image(self, img: np.ndarray):
 
     def get_points(self):
         """Get all selected points from manual and polygons."""
-        return np.array(self.selected_points)
+        points = np.array(self.selected_points)[:, ::-1] # set axis 0 to y and axis 1 to x
+        return points
+    
+    @property
+    def points(self):
+        return self.get_points()
 
     def get_filtered_points(self):
         """Get candidate points from automatic filtering."""