diff --git a/src/hflow/episode.py b/src/hflow/episode.py index 458faf45..6887f1e0 100644 --- a/src/hflow/episode.py +++ b/src/hflow/episode.py @@ -88,8 +88,28 @@ def _sanitize_topic(topic: str) -> str: return f"{readable}-{digest}" -def _frame_selection_expression(frame_indices: Sequence[int]) -> str: - """Build a compact ffmpeg select expression for sorted, unique indices.""" +# The cache directory name is the cache key, so the layout tag retires every +# directory written before the selection moved to timestamps (#727): those hold +# the decoded-counter pictures, which are shifted after a lost keyframe. +_FRAMES_AT_INDICES_CACHE_LAYOUT = "v2-time" + + +def _frame_selection_expression(frame_indices: Sequence[int], *, fps: float) -> str: + """Build a compact ffmpeg select expression for sorted, unique indices. + + Pictures are selected by time, not by ffmpeg's decoded-picture counter + ``n``: a decoder that loses a reference frame drops every picture that + depends on it, so decoded picture *n* stops being message *n* after a + missing keyframe and each later index returns a picture shifted by the + shortfall, labelled with the requested message's log time (#727). The MP4 + is constant rate, so message *k* is the picture nearest ``k / fps``, and a + picture the decoder never produced matches nothing -- the caller's + extracted-count check then reports the request instead of returning a + shifted frame. + """ + + if fps <= 0: + raise ValueError(f"fps must be positive, got {fps}") consecutive_ranges: list[tuple[int, int]] = [] range_start = frame_indices[0] @@ -102,10 +122,15 @@ def _frame_selection_expression(frame_indices: Sequence[int]) -> str: range_start = frame_index range_end = frame_index consecutive_ranges.append((range_start, range_end)) + + def half_frame_before(index: int) -> str: + return f"{(index - 0.5) / fps:.17g}" + + def half_frame_after(index: int) -> str: + return f"{(index + 0.5) / fps:.17g}" + return "+".join( - f"eq(n\\,{range_start})" - if range_start == range_end - else f"between(n\\,{range_start}\\,{range_end})" + f"between(t\\,{half_frame_before(range_start)}\\,{half_frame_after(range_end)})" for range_start, range_end in consecutive_ranges ) @@ -837,7 +862,7 @@ def frames_at_indices( selection_digest = hashlib.sha256(serialized_indices.encode()).hexdigest()[:16] output_directory = self.workdir / ( f"frames_{_sanitize_topic(topic)}_indices_{selection_digest}_" - f"count_{len(selected_frame_indices)}" + f"count_{len(selected_frame_indices)}_{_FRAMES_AT_INDICES_CACHE_LAYOUT}" ) expected_frame_paths = [ output_directory / f"frame_{output_index:06d}.jpg" @@ -852,7 +877,9 @@ def frames_at_indices( ) try: filter_script_path = staging_directory / "select.filter" - selection_expression = _frame_selection_expression(selected_frame_indices) + selection_expression = _frame_selection_expression( + selected_frame_indices, fps=self._video_fps[topic] + ) filter_script_path.write_text(f"select={selection_expression}") command = [ str(ffmpeg_path()), diff --git a/tests/test_end_to_end.py b/tests/test_end_to_end.py index 21ac53c5..47ac811a 100644 --- a/tests/test_end_to_end.py +++ b/tests/test_end_to_end.py @@ -2,20 +2,25 @@ infrastructure. Mirrors the README design-target example.""" import asyncio +import hashlib import logging import subprocess +from fractions import Fraction from pathlib import Path from tempfile import TemporaryDirectory from typing import Any +import av import numpy as np import pytest +from mcap_test_helpers import write_compressed_video_mcap import hflow from hflow.checks import camera_frame_stats from hflow.curation import open_catalog_connection from hflow.ffmpeg._binary import _ffmpeg_filter_script_flag from hflow.testing import SyntheticEpisodeSpec, synthesize_episode +from hflow.transform import TransformConfig, write_canonical_episode def check_joint_smoothness(joints: np.ndarray, rate_hz: float) -> dict[str, float]: @@ -520,6 +525,192 @@ def refuse_cached_read(*args: Any, **kwargs: Any) -> None: ) +@pytest.mark.requires_system_ffmpeg +def test_frames_at_indices_pairs_pictures_by_time_when_a_keyframe_is_missing( + tmp_path: Path, +) -> None: + """#727: a lost keyframe must not shift the picture of every later index. + + The picture carries its own source frame number, so a returned JPEG is + checked against the message that was asked for instead of against a count. + """ + + frame_count, frame_rate, gop_frames = 60, 30, 10 + lost_source_frame = 30 + topic = "/camera/cam0/video" + first_log_time_ns = 1_760_000_000_000_000_000 + access_units = _encode_index_pattern_stream( + frame_count=frame_count, frame_rate=frame_rate, gop_frames=gop_frames + ) + assert len(access_units) == frame_count + + kept_source_frames = [index for index in range(frame_count) if index != lost_source_frame] + raw_path = tmp_path / "raw.mcap" + write_compressed_video_mcap( + raw_path, + [ + ( + topic, + first_log_time_ns + source_frame * 1_000_000_000 // frame_rate, + access_units[source_frame], + ) + for source_frame in kept_source_frames + ], + ) + canonical_path = tmp_path / "canonical.mcap" + write_canonical_episode(str(raw_path), str(canonical_path), TransformConfig()) + + with hflow.Episode(str(canonical_path)) as episode: + after_the_gap = [41, 50] + extracted = episode.frames_at_indices(topic, frame_indices=after_the_gap) + assert [_read_frame_number(frame.path) for frame in extracted] == [ + kept_source_frames[index] for index in after_the_gap + ] + + inside_the_gap = [33, 34] + with pytest.raises(RuntimeError, match="extracted"): + episode.frames_at_indices(topic, frame_indices=inside_the_gap) + + +@pytest.mark.requires_system_ffmpeg +def test_frames_at_indices_re_extracts_a_cache_written_by_the_counter_selection( + tmp_path: Path, +) -> None: + """#727 upgrade: a workdir populated before the fix must not be reused. + + ``main`` cached its counter selection under + ``frames__indices__count_``, so a workdir built there + holds the shifted pictures. Reopening it on this head has to extract again + instead of serving them. + """ + + frame_count, frame_rate, gop_frames = 60, 30, 10 + lost_source_frame = 30 + topic = "/camera/cam0/video" + first_log_time_ns = 1_760_000_000_000_000_000 + access_units = _encode_index_pattern_stream( + frame_count=frame_count, frame_rate=frame_rate, gop_frames=gop_frames + ) + kept_source_frames = [index for index in range(frame_count) if index != lost_source_frame] + raw_path = tmp_path / "raw.mcap" + write_compressed_video_mcap( + raw_path, + [ + ( + topic, + first_log_time_ns + source_frame * 1_000_000_000 // frame_rate, + access_units[source_frame], + ) + for source_frame in kept_source_frames + ], + ) + canonical_path = tmp_path / "canonical.mcap" + write_canonical_episode(str(raw_path), str(canonical_path), TransformConfig()) + + requested = [41, 50] + workdir = tmp_path / "workdir" + stale_cache = workdir / _counter_selection_cache_directory(topic, requested) + stale_cache.mkdir(parents=True) + for output_index, shifted_source_frame in enumerate((50, 59)): + _write_pattern_jpeg(stale_cache / f"frame_{output_index:06d}.jpg", shifted_source_frame) + + with hflow.Episode(str(canonical_path), workdir=workdir) as episode: + extracted = episode.frames_at_indices(topic, frame_indices=requested) + + assert [_read_frame_number(frame.path) for frame in extracted] == [ + kept_source_frames[index] for index in requested + ] + assert not any(frame.path.parent == stale_cache for frame in extracted) + + +def _counter_selection_cache_directory(topic: str, frame_indices: list[int]) -> str: + """The frame-cache directory ``main`` wrote while it selected by decoded picture.""" + + serialized = ",".join(str(frame_index) for frame_index in frame_indices) + digest = hashlib.sha256(serialized.encode()).hexdigest()[:16] + return ( + f"frames_{hflow.episode._sanitize_topic(topic)}_indices_{digest}_count_{len(frame_indices)}" + ) + + +def _write_pattern_jpeg(jpeg_path: Path, source_frame: int) -> None: + """Write one frame-number picture as a JPEG, the shape a frame cache holds.""" + + with av.open(str(jpeg_path), mode="w", format="image2") as container: + stream = container.add_stream("mjpeg", rate=1) + stream.width, stream.height = _PATTERN_WIDTH, _PATTERN_HEIGHT + stream.pix_fmt = "yuvj420p" + for packet in stream.encode( + av.VideoFrame.from_ndarray(_index_pattern_picture(source_frame), format="rgb24") + ): + container.mux(packet) + for packet in stream.encode(None): + container.mux(packet) + + +_PATTERN_BITS = 10 +_PATTERN_WIDTH = 320 +_PATTERN_HEIGHT = 180 +_PATTERN_BAND_ROWS = 32 + + +def _index_pattern_picture(source_frame: int) -> np.ndarray: + """A gray picture whose top band is the source frame number in black/white blocks.""" + + picture = np.full((_PATTERN_HEIGHT, _PATTERN_WIDTH, 3), 128, np.uint8) + block_width = _PATTERN_WIDTH // _PATTERN_BITS + picture[:_PATTERN_BAND_ROWS] = 16 + for bit in range(_PATTERN_BITS): + if (source_frame >> (_PATTERN_BITS - 1 - bit)) & 1: + left = bit * block_width + picture[:_PATTERN_BAND_ROWS, left : left + block_width] = 235 + return picture + + +def _encode_index_pattern_stream( + *, frame_count: int, frame_rate: int, gop_frames: int +) -> list[bytes]: + """Encode one access unit per frame, each carrying its own source frame number.""" + + context = av.CodecContext.create("libx264", "w") + context.width, context.height, context.pix_fmt = ( + _PATTERN_WIDTH, + _PATTERN_HEIGHT, + "yuv420p", + ) + context.framerate, context.time_base = Fraction(frame_rate), Fraction(1, frame_rate) + context.gop_size, context.max_b_frames = gop_frames, 0 + context.options = { + "profile": "baseline", + "tune": "zerolatency", + "keyint_min": str(gop_frames), + "sc_threshold": "0", + } + access_units: list[bytes] = [] + for source_frame in range(frame_count): + picture = av.VideoFrame.from_ndarray( + _index_pattern_picture(source_frame), format="rgb24" + ).reformat(format="yuv420p") + picture.pts = source_frame + access_units.extend(bytes(packet) for packet in context.encode(picture)) + access_units.extend(bytes(packet) for packet in context.encode(None)) + return access_units + + +def _read_frame_number(jpeg_path: Path) -> int: + """Read the source frame number back out of an extracted JPEG.""" + + with av.open(str(jpeg_path)) as container: + picture = next(container.decode(video=0)).to_ndarray(format="gray") + band = picture[:_PATTERN_BAND_ROWS] + block_width = _PATTERN_WIDTH // _PATTERN_BITS + bits = "".join( + "1" if band[:, bit * block_width : (bit + 1) * block_width].mean() > 128 else "0" + for bit in range(_PATTERN_BITS) + ) + return int(bits, 2) + + def test_canonical_episode_sparse_selection_past_argument_limit_uses_filter_file( report_and_app: tuple[hflow.TestReport, hflow.App], monkeypatch: pytest.MonkeyPatch, @@ -531,7 +722,12 @@ def test_canonical_episode_sparse_selection_past_argument_limit_uses_filter_file """ report, _app = report_and_app sparse_indices = list(range(0, 40000, 2)) - expression = hflow.episode._frame_selection_expression(sparse_indices) + with hflow.Episode(report.canonical_path) as fps_source: + fps_source.video(next(topic for topic in fps_source.cameras if "overhead_cam" in topic)) + source_fps = fps_source._video_fps[ + next(topic for topic in fps_source.cameras if "overhead_cam" in topic) + ] + expression = hflow.episode._frame_selection_expression(sparse_indices, fps=source_fps) assert len(expression) > 131072 # Exceeds Linux 128 KiB limit with hflow.Episode(report.canonical_path) as episode: diff --git a/tests/test_episode.py b/tests/test_episode.py index e9c05eeb..c5335775 100644 --- a/tests/test_episode.py +++ b/tests/test_episode.py @@ -1,10 +1,11 @@ +import re from pathlib import Path from typing import Any import pytest from mcap.writer import CompressionType, Writer -from hflow.episode import Episode +from hflow.episode import Episode, _frame_selection_expression def _write_two_json_channels( @@ -309,3 +310,22 @@ def test_channel_to_arrow_numeric_list_with_null_elements() -> None: [None, None, None], [1.0, 2.0, 3.0], ] + + +def test_frame_selection_expression_keeps_large_index_windows_one_frame_wide() -> None: + """Six-digit time windows must not swallow the neighbouring picture. + + ``%g`` prints six significant digits, so index 250001 at 25 fps collapsed to + ``between(t,10000,10000.1)`` -- a window covering indices 250000-250002. The + boundaries must carry round-trip precision instead (#730 review). + """ + fps = 25.0 + frame_indices = [250001, 250010] + expression = _frame_selection_expression(frame_indices, fps=fps) + + windows = re.findall(r"between\(t\\,([^,\\]+)\\,([^)\\]+)\)", expression) + assert len(windows) == len(frame_indices) + for frame_index, (start, end) in zip(frame_indices, windows, strict=True): + start_time, end_time = float(start), float(end) + assert start_time < frame_index / fps < end_time + assert end_time - start_time == pytest.approx(1 / fps)