Add a grayscale flag and mandate frame_index in plot_frame()

Author: qian-chu

pyneon/utils/variables.py

@@ -115,11 +115,12 @@
     # Events columns after concatenation
     "message name": "string",
     "message type": "string",
-    # Marker detections
+    # Detections
     "frame index": "Int64",
     "marker family": "string",
     "marker id": "string",
     "marker name": "string",
+    "contour name": "string",
     "top left x [px]": float,
     "top left y [px]": float,
     "top right x [px]": float,

pyneon/video/video.py

@@ -457,7 +457,8 @@ def __exit__(self, exc_type, exc, traceback) -> bool:
     @fill_doc
     def plot_frame(
         self,
-        frame_index: int = 0,
+        frame_index: int,
+        grayscale: bool = False,
         ax: Optional[plt.Axes] = None,
         show: bool = True,
     ):
@@ -468,13 +469,17 @@ def plot_frame(
         ----------
         frame_index : int
             Index of the frame to plot.
+        grayscale : bool, optional
+            Whether to convert the frame to grayscale before plotting.
+            Defaults to False.
         {ax_param}
+        {show_param}
 
         Returns
         -------
         {fig_ax_returns}
         """
-        return plot_frame(self, frame_index, ax, show)
+        return plot_frame(self, frame_index, grayscale, ax, show)
 
     def undistort_video(
         self,

pyneon/vis/video.py

@@ -20,7 +20,8 @@
 @fill_doc
 def plot_frame(
     video: "Video",
-    frame_index: int = 0,
+    frame_index: int,
+    grayscale: bool = False,
     ax: Optional[plt.Axes] = None,
     show: bool = True,
 ) -> tuple[plt.Figure, plt.Axes]:
@@ -33,6 +34,9 @@ def plot_frame(
         Video instance to plot the frame from.
     frame_index : int
         Index of the frame to plot.
+    grayscale : bool, optional
+        Whether to convert the frame to grayscale before plotting.
+        Defaults to False.
     {ax_param}
     {show_param}
 
@@ -49,7 +53,11 @@ def plot_frame(
     frame = video.read_frame_at(frame_index)
     if frame is None:
         raise RuntimeError(f"Could not read frame {frame_index}")
-    ax.imshow(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+    if grayscale:
+        frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+        ax.imshow(frame, cmap="gray")
+    else:
+        ax.imshow(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
     ax.axis("off")
     if show:
         plt.show()