Make zarr time axis dense

crabs/zarr/create_dataset.py

@@ -66,8 +66,16 @@ def load_extended_ds(
         added metadata as coordinates and data variables.
 
     """
-    # Load VIA tracks file as movement dataset
-    ds = load_bboxes.from_via_tracks_file(via_tracks_file_path)
+    # Load VIA tracks file as movement dataset.
+    # We request the frame numbers as defined in the file. With the
+    # default (use_frame_numbers_from_file=False), `movement` maps them to
+    # a 0-based sequence of *consecutive* integers. Since the tracker
+    # writes no rows for frames with no tracked boxes, that would shift
+    # every subsequent frame earlier, and `time` would no longer be the
+    # clip's frame index.
+    ds = load_bboxes.from_via_tracks_file(
+        via_tracks_file_path, use_frame_numbers_from_file=True
+    )
 
     # Extract metadata for this row
     clip_filename = _via_tracks_to_clip_filename(ds.attrs["source_file"])
@@ -77,14 +85,43 @@ def load_extended_ds(
     global_clip_start_frame_0idx = row["loop_START_frame_ffmpeg"] - 1
     global_clip_end_frame_0idx = row["loop_END_frame_ffmpeg"] - 1
     global_escape_start_frame_0idx = row["escape_START_frame_0_based_idx"]
+    n_clip_frames = (
+        global_clip_end_frame_0idx - global_clip_start_frame_0idx + 1
+    )
+
+    # Check the frames in the VIA tracks file belong to this clip
+    _validate_frames_in_clip_range(
+        ds.time.values,
+        n_clip_frames,
+        global_clip_start_frame_0idx,
+        clip_filename,
+    )
+
+    # Pad the time coordinate to span the full clip.
+    # `movement` pads with NaNs the individuals that are missing from a
+    # frame, but a frame that no individual appears in is absent from the
+    # time coordinate altogether, so we reindex to add it back.
+    n_frames_wout_boxes = n_clip_frames - ds.sizes["time"]
+    if n_frames_wout_boxes:
+        print(
+            f"{clip_filename}: {n_frames_wout_boxes} of {n_clip_frames} "
+            "frames have no tracked boxes; padding them with NaNs."
+        )
+    ds = ds.reindex(time=np.arange(n_clip_frames))
 
     # Add escape_state as data variable
     local_escape_start_frame_0idx = (
         global_escape_start_frame_0idx - global_clip_start_frame_0idx
     )
+    if not 0 <= local_escape_start_frame_0idx < n_clip_frames:
+        raise ValueError(
+            f"{clip_filename}: the escape start frame "
+            f"({local_escape_start_frame_0idx} as a clip frame index) is "
+            f"outside the clip's frame range 0-{n_clip_frames - 1}."
+        )
     # we use float16 (not int/bool) to allow for NaN padding after
     # concatenating along clip_id
-    escape_state = np.zeros(ds.time.shape[0], dtype=np.float16)
+    escape_state = np.zeros(n_clip_frames, dtype=np.float16)
     escape_state[local_escape_start_frame_0idx:] = 1.0
     ds["escape_state"] = ("time", escape_state)
 
@@ -118,6 +155,33 @@ def load_extended_ds(
     return ds
 
 
+def _validate_frames_in_clip_range(
+    frames_0idx: np.ndarray,
+    n_clip_frames: int,
+    clip_first_frame_0idx: int,
+    clip_filename: str,
+) -> None:
+    """Check the frame numbers in a VIA tracks file belong to the clip.
+
+    The frame numbers are expected to be 0-based indices into the clip,
+    which is how `detect-and-track-video` numbers the frames of the clip
+    it is given.
+    """
+    if frames_0idx.min() >= 0 and frames_0idx.max() < n_clip_frames:
+        return
+
+    raise ValueError(
+        f"{clip_filename}: the VIA tracks file holds frame numbers "
+        f"{frames_0idx.min()}-{frames_0idx.max()}, outside the clip's "
+        f"frame range 0-{n_clip_frames - 1}. Frame numbers are expected "
+        f"to be 0-based indices into the clip, which starts at video "
+        f"frame {clip_first_frame_0idx}. Either the file was tracked on "
+        f"the full video or paired with the wrong metadata row, or the "
+        f"cut clip does not have the frame count the metadata csv states "
+        f"(re-run extract-loop-clips --verify_frames)."
+    )
+
+
 def _group_files_per_video(
     files_dir: str | Path,
     glob_pattern: str,

guides/CreateZarrDatasetForTracks.md

@@ -80,6 +80,21 @@
 
     If the array job runs successfully, a zarr dataset named `CrabTracks-slurm<SLURM_ARRAY_JOB_ID>.zarr` will be generated in the location specified by `ZARR_STORE_OUTPUT`. Each group in the zarr dataset will correspond to a `movement` [bounding box dataset](https://movement.neuroinformatics.dev/latest/user_guide/movement_dataset.html) containing all tracks for one video.
 
+    The `time` coordinate is the clip's 0-based frame index, and it always spans the whole clip: frames that hold no tracked crabs are stored as NaNs rather than dropped. So within a clip, row `i` is frame `i`.
+
+    To get a single clip's full frame range, without the NaN padding that the concatenation along `clip_id` adds beyond the end of shorter clips:
+
+    ```python
+    d = ds_video.isel(clip_id=0)
+    n_clip_frames = int(d.clip_last_frame_0idx - d.clip_first_frame_0idx) + 1
+    d_clip = d.isel(time=slice(0, n_clip_frames))
+    ```
+
+    Prefer this to `dropna(dim="time", how="all")`, which cannot tell "beyond this clip's end" from "no crabs in this frame": it silently drops crab-free frames from the middle of the clip and returns a gapped time axis. The slice keeps every frame of the clip and nothing else.
+
+    > [!WARNING]
+    > Zarr stores built before the `time` coordinate was anchored to the clip have a `time` coordinate that does not mean this. There, `time` was a 0-based count of the frames *with detections*, so frames with no tracked crabs shifted every later frame earlier, and `escape_state` switched on at the wrong frame. Those stores may need rebuilding.
+
     To inspect this dataset, please refer to the example notebooks in the 🦀 repository.
 
 ### Re-running failed jobs

scripts/visualization/output_video.py

@@ -5,6 +5,7 @@
 
 import cv2
 import matplotlib.pyplot as plt
+import numpy as np
 from movement.io import load_bboxes
 from tqdm import tqdm
 
@@ -34,6 +35,50 @@ def get_distinct_colors():
     return colors
 
 
+def reindex_to_video_frames(ds, video_path):
+    """Pad the dataset's time coordinate to span every frame of the video.
+
+    `movement` pads with NaNs the individuals that are missing from a frame,
+    but a frame that no individual appears in is absent from the time
+    coordinate altogether. We reindex to add those frames back, so that row
+    `i` of the dataset is frame `i` of the video and the arrays can be
+    indexed positionally by frame number.
+
+    The dataset is expected to hold the frame numbers as written in the VIA
+    tracks file (i.e. loaded with `use_frame_numbers_from_file=True` and
+    `fps=None`), which the tracker numbers from 0 over the video it is given.
+    """
+    # Get n frames in video
+    cap = cv2.VideoCapture(str(video_path))
+    n_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+    cap.release()
+
+    # Check the frame numbers belong to this video. `reindex` discards
+    # out-of-range labels silently, so without this an unrelated tracks file
+    # would give an all-NaN dataset rather than an error.
+    frames_0idx = ds.time.values
+    if frames_0idx.min() < 0 or frames_0idx.max() >= n_frames:
+        raise ValueError(
+            f"{Path(video_path).name}: the frame numbers in the tracks file "
+            f"({frames_0idx.min()}-{frames_0idx.max()}) fall outside the "
+            f"video's frame range 0-{n_frames - 1}. Are the tracks from this "
+            "video?"
+        )
+
+    # Print number of frames with no detections
+    n_frames_wo_boxes = n_frames - ds.sizes["time"]
+    if n_frames_wo_boxes:
+        print(
+            f"{Path(video_path).name}: {n_frames_wo_boxes} of {n_frames} "
+            "frames have no tracked boxes; padding them with NaNs."
+        )
+
+    # Reindex dataset to number of frames in video
+    # (reindexing fills in any gaps with nans, but silently drops
+    # frames out of the specificied range; hence the range check earlier)
+    return ds.reindex(time=np.arange(n_frames))
+
+
 def create_opencv_video(
     ds,
     input_video,
@@ -172,10 +217,19 @@ def create_opencv_video(
         input_video = Path(input_data_dir) / f"{escape_clip_name}.mp4"
         pred_csv = Path(input_data_dir) / f"{escape_clip_name}_tracks.csv"
 
-        # Read predictions as a movement dataset
+        # Read predictions as a movement dataset.
+        # We request the frame numbers as defined in the file
+        # (only frames with detections are included in the file)
         ds_pred = load_bboxes.from_via_tracks_file(
-            pred_csv, fps=None, use_frame_numbers_from_file=False
+            pred_csv,
+            fps=None,
+            use_frame_numbers_from_file=True,
         )
+
+        # Pad the time coordinate to span the full video, so that row i of
+        # the dataset is frame i of the video
+        ds_pred = reindex_to_video_frames(ds_pred, input_video)
+
         list_individuals_idcs = list(range(len(ds_pred.individuals)))
 
         # Create prediction video

tests/test_unit/test_create_zarr_dataset.py

@@ -2,6 +2,7 @@
 from contextlib import nullcontext as does_not_raise
 from pathlib import Path
 
+import numpy as np
 import pandas as pd
 import pytest
 import xarray as xr
@@ -24,11 +25,16 @@
 def sample_via_tracks_file_factory(tmp_path):
     """Return factory of VIA tracks files with custom filenames."""
 
-    def _sample_via_tracks_file(filename):
-        """Create a mock VIA tracks file."""
+    def _sample_via_tracks_file(filename, frame_indices=None):
+        """Create a mock VIA tracks file.
+
+        By default the file holds boxes for a sparse set of frames of a
+        101-frame clip, which is the case the dense `time` axis is about.
+        Pass `frame_indices` to write boxes for a different set of frames.
+        """
         # Add a few frames with bounding box annotations
         via_data = []
-        for frame_idx in [0, 10, 20, 50, 60, 80, 100]:
+        for frame_idx in frame_indices or [0, 10, 20, 50, 60, 80, 100]:
             via_data.append(
                 {
                     "filename": f"frame_{frame_idx:08d}.png",
@@ -151,6 +157,96 @@ def test_load_extended_ds(
     assert "clip_escape_first_frame_0idx" in ds.coords
     assert "clip_escape_type" in ds.coords
 
+    # Check time spans the full clip, as the clip's 0-based frame index
+    assert ds.sizes["time"] == 101
+    assert np.array_equal(ds.time.values, np.arange(101))
+
+    # Check the escape starts at clip frame 51
+    # (global frame 150, clip starting at global frame 99)
+    escape_state = ds.escape_state.values.ravel()
+    assert int(np.argmax(escape_state)) == 51
+    assert escape_state.sum() == 101 - 51
+
+    # Check frames with no boxes are all-NaN rows, rather than missing
+    # (the VIA tracks file holds boxes for frames 0, 10, 20, 50, 60, 80, 100)
+    assert ds.position.isel(time=1).isnull().all()
+    assert not ds.position.isel(time=10).isnull().any()
+
+
+def test_load_extended_ds_contiguous_file(
+    sample_via_tracks_file_factory, sample_metadata_df_factory
+):
+    """Test a VIA tracks file with one row per frame of the clip.
+
+    It should give the same `time` and `escape_state` as the sparse file,
+    and no NaN rows.
+    """
+    via_tracks_path = sample_via_tracks_file_factory(
+        "04.09.2023-01-Right-Loop05_tracks.csv",
+        frame_indices=list(range(101)),
+    )
+    df_metadata = sample_metadata_df_factory([via_tracks_path])
+
+    ds = load_extended_ds(via_tracks_path, df_metadata)
+
+    assert np.array_equal(ds.time.values, np.arange(101))
+    escape_state = ds.escape_state.values.ravel()
+    assert int(np.argmax(escape_state)) == 51
+    assert escape_state.sum() == 101 - 51
+    assert not ds.position.isnull().any()
+
+
+@pytest.mark.parametrize(
+    "frame_indices",
+    [
+        pytest.param(
+            list(range(99, 200)),
+            id="tracked_on_full_video_or_wrong_metadata_row",
+        ),
+        pytest.param(
+            list(range(102)), id="cut_clip_longer_than_metadata_states"
+        ),
+    ],
+)
+def test_load_extended_ds_frames_out_of_clip_range(
+    frame_indices, sample_via_tracks_file_factory, sample_metadata_df_factory
+):
+    """Test the guard on frame numbers outside the clip's frame range.
+
+    Without it, `reindex` would silently drop the out-of-range rows.
+    """
+    via_tracks_path = sample_via_tracks_file_factory(
+        "04.09.2023-01-Right-Loop05_tracks.csv", frame_indices=frame_indices
+    )
+    df_metadata = sample_metadata_df_factory([via_tracks_path])
+
+    with pytest.raises(ValueError) as exc_info:
+        load_extended_ds(via_tracks_path, df_metadata)
+
+    # The message reports the file's own range and the clip's range
+    message = str(exc_info.value)
+    assert f"{frame_indices[0]}-{frame_indices[-1]}" in message
+    assert "0-100" in message
+
+
+def test_load_extended_ds_escape_frame_out_of_clip_range(
+    sample_via_tracks_file_factory, sample_metadata_df_factory
+):
+    """Test the guard on an escape frame outside the clip's frame range."""
+    via_tracks_path = sample_via_tracks_file_factory(
+        "04.09.2023-01-Right-Loop05_tracks.csv"
+    )
+    df_metadata = sample_metadata_df_factory([via_tracks_path])
+
+    # Clip spans global frames 99-199; put the escape past its end
+    df_metadata.loc[0, "escape_START_frame_0_based_idx"] = 250
+
+    with pytest.raises(ValueError) as exc_info:
+        load_extended_ds(via_tracks_path, df_metadata)
+
+    assert "escape start frame" in str(exc_info.value)
+    assert "0-100" in str(exc_info.value)
+
 
 @pytest.mark.parametrize(
     "list_files, expected_exception",