diff --git a/crabs/zarr/create_dataset.py b/crabs/zarr/create_dataset.py index fa02c464..6e519241 100644 --- a/crabs/zarr/create_dataset.py +++ b/crabs/zarr/create_dataset.py @@ -66,8 +66,12 @@ 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. Note that + # the tracker only writes frames with tracked boxes. + 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 +81,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 are all + # within range + _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. + 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." + ) + # Reindex time coordinates to fill in empy frames with NaN + # (out-of-range frames are silently dropped, hence the range check above) + 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 index " + f"({local_escape_start_frame_0idx} as a clip frame index) is " + f"outside the clip's 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 +151,31 @@ 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 are within the clip. + + The frame indices are 0-based indices into the clip. + """ + if frames_0idx.min() >= 0 and frames_0idx.max() < n_clip_frames: + return + + raise ValueError( + f"{clip_filename}: the VIA tracks file holds frame indices " + f"{frames_0idx.min()}-{frames_0idx.max()}, outside the clip's " + f"range 0-{n_clip_frames - 1}. Frame indices are expected " + f"to be 0-based indices into the clip, which starts at video " + f"frame {clip_first_frame_0idx}. Either the frames in the file " + f"refer to the full video, the file is paired with the wrong " + f"metadata row, or the clip does not have the frame count the " + f"metadata csv states. " + ) + + def _group_files_per_video( files_dir: str | Path, glob_pattern: str, diff --git a/guides/CreateZarrDatasetForTracks.md b/guides/CreateZarrDatasetForTracks.md index 74e09a4e..3c53af7f 100644 --- a/guides/CreateZarrDatasetForTracks.md +++ b/guides/CreateZarrDatasetForTracks.md @@ -80,6 +80,21 @@ If the array job runs successfully, a zarr dataset named `CrabTracks-slurm.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. So within a clip, index `i` is frame `i`. + + To get a single clip's full frame range from a video dataset, without the trailing NaN padding that the concatenation along `clip_id` adds: + + ```python + ds = ds_video.isel(clip_id=0) + n_clip_frames = int(ds.clip_last_frame_0idx - ds.clip_first_frame_0idx) + 1 + ds_clip = ds.isel(time=slice(0, n_clip_frames)) # or equivalently: ds.sel(time=slice(0, n_clip_frames - 1)). Note that .sel is end-inclusive. + ``` + + Notice that when we drop all NaN frames with `dropna(dim="time", how="all")`, we cannot tell which NaN frames are padding and which NaN frames are "no detections in this frame": it silently returns a potentially gapped time axis. Instead the slice keeps every frame of the clip. + + > [!WARNING] + > Zarr stores built before the `time` coordinate was made dense 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 if there were frames with no detections in them. + To inspect this dataset, please refer to the example notebooks in the 🦀 repository. ### Re-running failed jobs diff --git a/scripts/visualization/output_video.py b/scripts/visualization/output_video.py index 7199d908..f42d0100 100644 --- a/scripts/visualization/output_video.py +++ b/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,42 @@ def get_distinct_colors(): return colors +def reindex_to_video_clip_frames(ds, video_path): + """Pad the dataset's time coordinate to span every frame of the clip. + + The dataset is expected to hold the frame numbers 0-based indices over + the clip, as written in the VIA tracks file (i.e. loaded with + `use_frame_numbers_from_file=True` and `fps=None`). + """ + # Get n frames in video clip + cap = cv2.VideoCapture(str(video_path)) + n_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + cap.release() + + # Check the frame indices in the file are within the video range + 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 " + "clip?" + ) + + # 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, @@ -169,13 +206,22 @@ def create_opencv_video( ] ) for escape_clip_name in list_escape_clips[-1:]: - input_video = Path(input_data_dir) / f"{escape_clip_name}.mp4" + input_clip = 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 time + # coordinate i of the dataset is frame index (0-based) i of the clip + ds_pred = reindex_to_video_clip_frames(ds_pred, input_clip) + list_individuals_idcs = list(range(len(ds_pred.individuals))) # Create prediction video @@ -184,7 +230,7 @@ def create_opencv_video( ) create_opencv_video( ds=ds_pred, - input_video=input_video, + input_video=input_clip, output_video_path=output_video_path, list_individuals_idcs=list_individuals_idcs, ) diff --git a/tests/fixtures/integration.py b/tests/fixtures/integration.py index 98439557..7c1a8972 100644 --- a/tests/fixtures/integration.py +++ b/tests/fixtures/integration.py @@ -5,7 +5,7 @@ import pooch import pytest -GIN_TEST_DATA_REPO = "https://gin.g-node.org/SainsburyWellcomeCentre/crabs-exploration-test-data" +GIN_TEST_DATA_REPO = "https://gin.swc.ucl.ac.uk/SainsburyWellcomeCentre/crabs-exploration-test-data" @pytest.fixture(scope="session") diff --git a/tests/test_unit/test_create_zarr_dataset.py b/tests/test_unit/test_create_zarr_dataset.py index 8b3a75d3..c4630616 100644 --- a/tests/test_unit/test_create_zarr_dataset.py +++ b/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 @@ -19,16 +20,25 @@ load_extended_ds, ) +# Frames the sample VIA tracks file holds boxes for by default +DEFAULT_FRAMES_WITH_BOXES = [0, 10, 20, 50, 60, 80, 100] + @pytest.fixture 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 few frames in a 101-frame + clip ("sparse" case). Pass `frame_indices` to write boxes + for a different set of frames, e.g. `range(101)` for the "dense" + case of one row per frame of the clip. + """ # 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 DEFAULT_FRAMES_WITH_BOXES: via_data.append( { "filename": f"frame_{frame_idx:08d}.png", @@ -122,21 +132,41 @@ def test_group_files_per_video(tmp_path): assert len(map_videos_to_files["04.09.2023-02-Right"]) == 1 +@pytest.mark.parametrize( + "via_frame_indices, expected_frames_in_ds", + [ + ( + None, # use default sparse frames from fixture + DEFAULT_FRAMES_WITH_BOXES, + ), + ( + list(range(101)), + list(range(101)), + ), + ], + ids=[ + "sparse_file", + "dense_file", + ], +) def test_load_extended_ds( - sample_via_tracks_file_factory, sample_metadata_df_factory + via_frame_indices, + expected_frames_in_ds, + sample_via_tracks_file_factory, + sample_metadata_df_factory, ): - """Test loading VIA tracks with metadata as a `movement` dataset.""" - # Sample VIA tracks file + """Test loading VIA tracks with metadata as an extended dataset.""" + # Get sample VIA tracks file via_tracks_filename = "04.09.2023-01-Right-Loop05_tracks.csv" - via_tracks_path = sample_via_tracks_file_factory(via_tracks_filename) + via_tracks_path = sample_via_tracks_file_factory( + via_tracks_filename, + frame_indices=via_frame_indices, + ) - # Define metadata df that contains the clip for the VIA tracks file - # list_clip_filepath = [ - # Path(str(via_tracks_path).replace("_tracks.csv", ".mp4")) - # ] + # Define metadata df for the VIA tracks file df_metadata = sample_metadata_df_factory([via_tracks_path]) - # Load dataset for this clip chunked + # Load extended dataset for this clip ds = load_extended_ds(via_tracks_path, df_metadata) # Check escape_state is a data var @@ -151,6 +181,91 @@ 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 + # the escape is at global frame 150 (0-based) and + # the clip starts at global frame 99 + # (loop_START_frame_ffmpeg=100, 1-indexed), + # so 150 - 99 = 51 + 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 in VIA file are all-NaN in the + # extended dataset + frames_not_nan = ( + ds.position.notnull() + .all(dim=["clip_id", "space", "individual"]) + .values + ) + assert np.array_equal( + np.flatnonzero(frames_not_nan), + expected_frames_in_ds, + ) + + +def test_frames_in_via_out_of_clip_range( + sample_via_tracks_file_factory, + sample_metadata_df_factory, +): + """Test the guard on VIA frame numbers outside the clip's frame range. + + Frame numbers in the VIA file are expected to be 0-based indices into + the clip. + """ + # Get VIA tracks file with one frame more than the clip + # (the clip frame indices range 0 to 100) + via_frame_indices = list(range(102)) # clip-based indices + via_tracks_path = sample_via_tracks_file_factory( + "04.09.2023-01-Right-Loop05_tracks.csv", + frame_indices=via_frame_indices, + ) + # Get clip metadata; it assumes the clip spans global frames 99-199, + # (i.e. clip frames 0-100) + df_metadata = sample_metadata_df_factory([via_tracks_path]) + + # Check ValueError is thrown + with pytest.raises(ValueError) as exc_info: + load_extended_ds(via_tracks_path, df_metadata) + + # Check message reports the file's own range and the clip's range + message = str(exc_info.value) + assert f"{via_frame_indices[0]}-{via_frame_indices[-1]}" in message + assert "0-100" in message + + +def test_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. + + It should be checked in load_extended_ds + """ + # Get VIA tracks file and metadata + 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]) + + # Place escape past the clip end. + # escape_START_frame_0_based_idx is in global (whole-video) numbering, + # and the clip spans global frames 99-199, so we place it at + # global frame 250 + df_metadata.loc[0, "escape_START_frame_0_based_idx"] = 250 + + # Check ValueError is thrown + with pytest.raises(ValueError) as exc_info: + load_extended_ds(via_tracks_path, df_metadata) + + # Check error message reports the clip-local index and range. + message = str(exc_info.value) + assert "escape start frame" in message + assert "151 as a clip frame index" in message + assert "0-100" in message + @pytest.mark.parametrize( "list_files, expected_exception",