Fix bug when overlaying contour detections on video, and improve tests/documentation

.github/workflows/main.yml

@@ -2,7 +2,7 @@ name: PyNeon CI
 
 on:
   push:
-    branches: ["main", "dev"]
+    branches: ["main"]
   pull_request:
     branches: ["main", "dev"]
 

pyneon/video/homography.py

@@ -8,18 +8,6 @@
 from .utils import _validate_contour_layout, _validate_marker_layout
 
 
-def _reshape_corners(detection: pd.Series) -> np.ndarray:
-    return np.array(
-        [
-            [detection["top left x [px]"], detection["top left y [px]"]],
-            [detection["top right x [px]"], detection["top right y [px]"]],
-            [detection["bottom right x [px]"], detection["bottom right y [px]"]],
-            [detection["bottom left x [px]"], detection["bottom left y [px]"]],
-        ],
-        dtype=np.float32,
-    )
-
-
 @fill_doc
 def find_homographies(
     detections: Stream,
@@ -46,7 +34,7 @@ def find_homographies(
         **Marker detections**: provide a DataFrame (can be visually checked with
         :func:`pyneon.plot_marker_layout`) with following columns:
 
-            {marker_layout_table}
+        {marker_layout_table}
 
         **Contour detections**: provide a 2D numpy array of shape (4, 2)
         containing the surface coordinates of the contour corners in the following order:
@@ -82,12 +70,12 @@ def find_homographies(
     Compute homographies from marker detections:
 
     >>> detections = video.detect_markers("36h11")
-    >>> layout = pd.DataFrame({
+    >>> layout = pd.DataFrame({{
     ...     "marker name": ["36h11_0", "36h11_1"],
     ...     "size": [100, 100],
     ...     "center x": [200, 400],
     ...     "center y": [200, 200],
-    ... })
+    ... }})
     >>> homographies = find_homographies(detections, layout)
 
     Compute homographies from contour detections:
@@ -103,57 +91,81 @@ def find_homographies(
     if is_marker_detection:
         # Validate marker layout
         _validate_marker_layout(layout)
-        # Compute corner coordinates for each marker in the layout
+        # Compute corner coordinates for each marker in the layout using vectorized arrays.
         layout = layout.copy()
-        layout["corners"] = layout.apply(
-            lambda row: np.array(
-                [
-                    [
-                        row["center x"] - row["size"] / 2,
-                        row["center y"] - row["size"] / 2,
-                    ],
-                    [
-                        row["center x"] + row["size"] / 2,
-                        row["center y"] - row["size"] / 2,
-                    ],
-                    [
-                        row["center x"] + row["size"] / 2,
-                        row["center y"] + row["size"] / 2,
-                    ],
-                    [
-                        row["center x"] - row["size"] / 2,
-                        row["center y"] + row["size"] / 2,
-                    ],
-                ],
-                dtype=np.float32,
-            ),
+        center_x = layout["center x"].to_numpy()
+        center_y = layout["center y"].to_numpy()
+        half_size = layout["size"].to_numpy() / 2
+        corners_array = np.stack(
+            [
+                np.column_stack((center_x - half_size, center_y - half_size)),
+                np.column_stack((center_x + half_size, center_y - half_size)),
+                np.column_stack((center_x + half_size, center_y + half_size)),
+                np.column_stack((center_x - half_size, center_y + half_size)),
+            ],
             axis=1,
         )
+        layout["corners"] = list(corners_array)
         # Construct a lookup dictionary with marker name being key and corners being value
         surface_pts_lookup = {
-            row["marker name"]: row["corners"] for _, row in layout.iterrows()
+            marker_name: corners
+            for marker_name, corners in layout[["marker name", "corners"]].itertuples(
+                index=False, name=None
+            )
         }
     else:
         _validate_contour_layout(layout)
         surface_pts_lookup = {"contour_0": layout}
 
+    corner_columns = [
+        "top left x [px]",
+        "top left y [px]",
+        "top right x [px]",
+        "top right y [px]",
+        "bottom right x [px]",
+        "bottom right y [px]",
+        "bottom left x [px]",
+        "bottom left y [px]",
+    ]
+
     homography_per_frame = {}
-    unique_timestamps = detection_df.index.unique()
+    grouped_detections = detection_df.groupby(level=0, sort=False)
 
-    for ts in tqdm(unique_timestamps, desc="Computing surface-mapping homographies"):
-        frame_detections = detection_df.loc[ts]
-        if isinstance(frame_detections, pd.Series):
-            frame_detections = frame_detections.to_frame().T
+    for ts, frame_detections in tqdm(
+        grouped_detections,
+        total=grouped_detections.ngroups,
+        desc="Computing surface-mapping homographies",
+    ):
+        if isinstance(frame_detections.index, pd.MultiIndex):
+            frame_detections = frame_detections.droplevel(0)
 
         if is_marker_detection and len(frame_detections) < min_markers:
             continue
 
         camera_pts_all = []
         surface_pts_all = []
 
-        for _, detection in frame_detections.iterrows():
-            camera_pts = _reshape_corners(detection)
-            name = detection["marker name"] if is_marker_detection else "contour_0"
+        corner_rows = frame_detections[corner_columns].itertuples(
+            index=False, name=None
+        )
+        if is_marker_detection:
+            detection_iter = zip(
+                frame_detections["marker name"].to_numpy(),
+                corner_rows,
+            )
+        else:
+            detection_iter = (("contour_0", corners) for corners in corner_rows)
+
+        for name, corners in detection_iter:
+            camera_pts = np.array(
+                [
+                    [corners[0], corners[1]],
+                    [corners[2], corners[3]],
+                    [corners[4], corners[5]],
+                    [corners[6], corners[7]],
+                ],
+                dtype=np.float32,
+            )
             surface_pts = surface_pts_lookup[name]
             if camera_pts.shape != (4, 2):
                 raise ValueError(

pyneon/vis/video.py

@@ -65,7 +65,7 @@ def plot_frame(
     return fig, ax
 
 
-def _overlay_marker_detections_on_frame(
+def _overlay_detections_on_frame(
     frame: np.ndarray,
     frame_detections: pd.DataFrame,
     show_ids: bool,
@@ -94,12 +94,17 @@ def _overlay_marker_detections_on_frame(
             [detection["center x [px]"], detection["center y [px]"]],
             dtype=np.int32,
         ).reshape(-1)
-        label = str(int(detection["marker id"])) if show_ids else None
 
         cv2.polylines(frame, [corners], True, color, 2)
 
-        if label is not None:
-            text = label
+        text = None
+        if show_ids:
+            if "marker id" in detection.index:
+                text = str(detection["marker id"])
+            elif "contour name" in detection.index:
+                text = str(detection["contour name"])
+
+        if text is not None:
             font = cv2.FONT_HERSHEY_SIMPLEX
             font_scale = 0.6
             thickness = 2
@@ -248,10 +253,13 @@ def overlay_detections(
         # Read the next frame sequentially
         frame = video.read_frame_at(frame_index)
 
+        if frame is None:  # Replace frame with all gray if it cannot be read
+            frame = np.full((video.height, video.width, 3), 128, dtype=np.uint8)
+
         if frame_index in detections_by_frame:
             frame_detections = detections_by_frame[frame_index]
             # Draw each detection on the frame
-            frame = _overlay_marker_detections_on_frame(
+            frame = _overlay_detections_on_frame(
                 frame, frame_detections, show_ids, color
             )
 
@@ -346,7 +354,7 @@ def overlay_scanpath(
         # Read the current frame from the video
         ret, frame = video.read()
         if not ret:
-            raise RuntimeError(f"Failed to read frame {idx} from the video.")
+            frame = np.full((video.height, video.width, 3), 128, dtype=np.uint8)
 
         # Extract fixations and gaze data
         fixations = row["fixations"]

source/reference/video.rst

@@ -1,13 +1,13 @@
-Video Class
-===========
+Video Class and Mapping Functions
+=================================
 
 This page documents the :class:`pyneon.Video` class, which provides an interface for working with video data in PyNeon.
-It also documents the :func:`pyneon.find_homographies` function, which computes homography transformations for video frames
-based on marker and contour layouts.
+It also documents the mapping related functions such as :func:`pyneon.find_homographies`.
 
 .. autoclass:: pyneon.Video
     :members:
     :show-inheritance:
 
 .. autofunction:: pyneon.find_homographies
 
+.. autofunction:: pyneon.plot_marker_layout

tests/conftest.py

@@ -237,3 +237,23 @@ def cloud_fixations(simple_dataset_cloud):
 @pytest.fixture(scope="package")
 def native_fixations(simple_dataset_native):
     return simple_dataset_native.recordings[0].fixations
+
+
+@pytest.fixture(scope="package")
+def mapping_dataset_native():
+    dataset_dir = get_sample_data("markers", format="native")
+    dataset = Dataset(dataset_dir)
+    assert len(dataset) == dataset.sections.shape[0] == 2
+    yield dataset
+    for recording in dataset.recordings:
+        recording.close()
+
+
+@pytest.fixture(scope="package")
+def mapping_dataset_cloud():
+    dataset_dir = get_sample_data("markers", format="cloud")
+    dataset = Dataset(dataset_dir)
+    assert len(dataset) == dataset.sections.shape[0] == 2
+    yield dataset
+    for recording in dataset.recordings:
+        recording.close()

tests/test_video.py

@@ -35,12 +35,16 @@ def test_video_basics(request, dataset_fixture):
             random_frames = np.append(random_frames, [0, n_frames - 1])
             for frame_idx in random_frames:
                 frame_idx = int(frame_idx)
-                frame = video.read_frame_at(frame_idx)
-                assert frame_idx == video.current_frame_index
                 if frame_idx == 0:
+                    with pytest.warns(
+                        UserWarning, match="Failed to retrieve frame at index 0"
+                    ):
+                        frame = video.read_frame_at(frame_idx)
                     assert frame is None
                 else:
+                    frame = video.read_frame_at(frame_idx)
                     assert frame.shape == (video.height, video.width, 3)
+                assert frame_idx == video.current_frame_index
 
             with pytest.raises(ValueError, match="is out of bounds."):
                 video.read_frame_at(-5)