[Bug]: rgb_frames(frames_per_second=…) samples come from after the message mediapipe_hands labels them with, and repeat across a lost keyframe #733

"""rgb_frames(frames_per_second=...): is each sampled picture the source frame its label names?

Labels are made the way mediapipe_hands makes them (mediapipe_hands.py:457-515): source_fps is one over the median
log-time interval, and source_log_times_for_sampled_frames maps sample i to a source message. MediaPipe itself is
not needed: the question is only which picture each label is attached to.

Synthetic data only. Each frame draws its source frame number as 10 black/white blocks, so the number can be read
back from the RGB array. One raw MCAP is complete, the other lacks source frame 100 (an IDR).

    pip install hflow av pillow mcap-protobuf-support foxglove-schemas-protobuf
    python hflow_rgb_frames_repro.py
"""
import tempfile
from fractions import Fraction
from pathlib import Path

import av
import numpy as np
from foxglove_schemas_protobuf.CompressedVideo_pb2 import CompressedVideo
from mcap.writer import CompressionType
from mcap_protobuf.writer import Writer

from hflow._video_measurement_toolchain import resolved_video_measurement_toolchain
from hflow._video_measurements._raw_frames import rgb_frames
from hflow.episode import Episode
from hflow.transform import TransformConfig, write_canonical_episode
from hflow.video import source_log_times_for_sampled_frames

W, H, FPS, N, GOP, BITS = 640, 360, 30, 300, 10, 10
T0 = 1_760_000_000_000_000_000
TOPIC = "/camera/cam0/video"


def draw(i):
    y = np.full((H, W), 128, np.uint8)
    y[:40] = 16
    for b in range(BITS):
        if (i >> (BITS - 1 - b)) & 1:
            y[:40, b * (W // BITS):(b + 1) * (W // BITS)] = 235
    return np.stack([y, y, y], -1)


def read_number(rgb):
    luma = rgb[:40].astype(np.float32).mean(axis=2)
    bw = W // BITS
    return int("".join(str(int(luma[:, b * bw + 8:(b + 1) * bw - 8].mean() > 128)) for b in range(BITS)), 2)


def encode():
    ctx = av.CodecContext.create("libx264", "w")
    ctx.width, ctx.height, ctx.pix_fmt = W, H, "yuv420p"
    ctx.framerate, ctx.time_base = FPS, Fraction(1, FPS)
    ctx.gop_size, ctx.max_b_frames = GOP, 0
    ctx.options = {"profile": "baseline", "tune": "zerolatency", "keyint_min": str(GOP), "sc_threshold": "0"}
    packets = []
    for i in range(N):
        frame = av.VideoFrame.from_ndarray(draw(i), format="rgb24").reformat(format="yuv420p")
        frame.pts = i
        packets += [bytes(p) for p in ctx.encode(frame)]
    return packets + [bytes(p) for p in ctx.encode(None)]


def write_raw(path, packets, drop):
    with open(path, "wb") as f:
        w = Writer(f, compression=CompressionType.ZSTD)
        for i in range(N):
            if i == drop:
                continue
            t = T0 + i * 1_000_000_000 // FPS
            msg = CompressedVideo(frame_id="cam0", data=packets[i], format="h264")
            msg.timestamp.FromNanoseconds(t)
            w.write_message(TOPIC, msg, log_time=t, publish_time=t)
        w.finish()


def source_frame_at(log_time_ns):
    return round((log_time_ns - T0) * FPS / 1e9)


def sample(ep, fps):
    """(source frame each label names, source frame each picture shows), as mediapipe_hands would label them."""
    stamps_ns = ep.channel(TOPIC).timestamps
    median_interval_s = float(np.median(np.diff(stamps_ns)) / 1e9)
    with rgb_frames(ep.video(TOPIC), toolchain=resolved_video_measurement_toolchain(),
                    frames_per_second=fps) as frames:
        shown = [read_number(f) for f in frames]
    labels = source_log_times_for_sampled_frames(stamps_ns.tolist(), source_fps=1.0 / median_interval_s,
                                                 sample_fps=fps, start_s=0.0, frame_count=len(shown))
    return list(zip([source_frame_at(t) for t in labels], shown))


def main():
    packets = encode()
    with tempfile.TemporaryDirectory() as work:
        for drop in (None, 100):
            raw, canonical = Path(work, f"raw_{drop}.mcap"), Path(work, f"canonical_{drop}.mcap")
            write_raw(raw, packets, drop)
            write_canonical_episode(str(raw), str(canonical), TransformConfig())
            print(f"\nsource {'complete' if drop is None else f'missing frame {drop} (an IDR)'}")
            ep = Episode(str(canonical), workdir=Path(work, f"ep_{drop}"))
            for fps in (30, 10, 3, 1):
                pairs = sample(ep, fps)
                late = sorted({shown - named for named, shown in pairs})
                print(f"  rgb_frames(frames_per_second={fps:2d}): {len(pairs):3d} frames; picture minus the frame "
                      f"its label names: {late}; first pairs (named, shown): {pairs[:3]}")
                if drop is not None:
                    gap = [p for p in pairs if drop - 1 <= p[0] <= drop + 10]
                    print(f"      samples named from {drop - 1} to {drop + 10}: {gap}")


if __name__ == "__main__":
    main()