VIDEO-005: add full offline QPSK IQ video receiver

tools/video005_video.py

@@ -0,0 +1,354 @@
+"""VIDEO-005: H.264 через непрерывный IQ, строгий PHY и внешний FEC."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass, replace
+from fractions import Fraction
+import hashlib
+import math
+import random
+import time
+
+import av
+from av.error import FFmpegError
+import numpy as np
+
+from protocol import punctured_fec
+from protocol.packet import parse_packet
+from protocol.video003_modem import PHY_HEADER_SYMBOLS, PREAMBLE_SYMBOLS, build_frame
+from protocol.video004_erasure import VideoErasureCollector, decode_source, parse_symbol
+from protocol.video005_iq_waveform import SAMPLE_RATE_HZ, rrc_taps, shape_packet
+from protocol.video005_receiver import IqStreamResult, receive_iq_stream
+from protocol.video_stream_transport import FrameReassembler, FrameType, VideoFrame
+from tools.video003_codec import SOURCE, encode_case
+from tools.video004_transport import _wire_schedule
+from tools.video_compare import source_frames
+from tools.video005_channel import impair
+
+
+STREAM_ID = 0x56303535
+FPS = 10
+SYMBOL_RATE = 500_000
+
+
+@dataclass(frozen=True)
+class Transmission:
+    index: int
+    start_sample: int
+    end_sample: int
+    block_id: int
+    source_index: int | None
+    source_frame_id: int | None
+    source_is_keyframe: bool
+
+
+@dataclass(frozen=True)
+class VideoFixture:
+    iq: np.ndarray
+    frames: tuple[VideoFrame, ...]
+    transmissions: tuple[Transmission, ...]
+    phy_packet_sha256: tuple[str, ...]
+    clean_pixel_sha256: dict[int, str]
+    source_sha256: str
+    encoder_settings: dict
+    preparation_time_s: dict[str, float]
+
+
+def _clean_pixels(packets: list[bytes]) -> dict[int, str]:
+    decoder = av.CodecContext.create("h264", "r")
+    decoder.open()
+    hashes = {}
+    for index, wire in enumerate(packets):
+        packet = av.Packet(wire)
+        packet.pts = packet.dts = index
+        packet.time_base = Fraction(1, FPS)
+        decoded = decoder.decode(packet)
+        if len(decoded) != 1 or decoded[0].pts != index:
+            raise AssertionError("исходный H.264 не декодировался точно")
+        hashes[index] = hashlib.sha256(decoded[0].to_ndarray(format="rgb24").tobytes()).hexdigest()
+    return hashes
+
+
+def build_fixture(frame_count: int = 20,
+                  pixels_override: list[np.ndarray] | None = None) -> VideoFixture:
+    """Кодирует исходный клип и передаёт видео по расписанию без oracle для RX."""
+    if not 1 <= frame_count <= 20:
+        raise ValueError("контроль VIDEO-005 ограничен первыми 20 кадрами")
+    started = time.perf_counter()
+    pixels = (source_frames(SOURCE, 640, 360, FPS, 0, frame_count / FPS)
+              if pixels_override is None else pixels_override)
+    if len(pixels) != frame_count or any(item.shape != (360, 640, 3) for item in pixels):
+        raise ValueError("ожидаются кадры RGB 640x360 выбранной длины")
+    codec, encoded = encode_case(pixels, FPS, "SHORT_GOP", 400_000)
+    if len(encoded) != frame_count:
+        raise AssertionError("кодек вернул неверное число кадров")
+    encode_time = time.perf_counter() - started
+    started = time.perf_counter()
+    clean_hashes = _clean_pixels(encoded)
+    reference_decode_time = time.perf_counter() - started
+    frames = tuple(VideoFrame(
+        STREAM_ID, index,
+        FrameType.KEYFRAME if index in codec["keyframe_positions"] else FrameType.DELTA_FRAME,
+        index * 100, wire,
+    ) for index, wire in enumerate(encoded))
+
+    started = time.perf_counter()
+    schedule, _ = _wire_schedule(list(frames), 8, 1, True, 1)
+    phy_frames = [build_frame(wire, "QPSK", "2/3", sequence_number=index)
+                  for index, (wire, _, _, _) in enumerate(schedule)]
+    outer_phy_time = time.perf_counter() - started
+
+    started = time.perf_counter()
+    chunks: list[np.ndarray] = []
+    intervals: list[Transmission] = []
+    position = 0
+    for index, ((wire, ready_ms, block_id, is_parity), phy) in enumerate(
+        zip(schedule, phy_frames, strict=True)
+    ):
+        ready_sample = math.ceil(ready_ms * SAMPLE_RATE_HZ / 1000)
+        if ready_sample > position:
+            chunks.append(np.zeros(ready_sample - position, dtype=np.complex64))
+            position = ready_sample
+        waveform = shape_packet(phy)
+        start_sample = position
+        chunks.append(waveform)
+        position += waveform.size
+        symbol = parse_symbol(wire)
+        source = (decode_source(symbol.data, symbol.stream_id, symbol.index, False)
+                  if not is_parity else None)
+        intervals.append(Transmission(
+            index, start_sample, position, block_id,
+            symbol.index if source is not None else None,
+            source.fragment.frame_id if source is not None else None,
+            bool(source is not None and source.fragment.frame_type == FrameType.KEYFRAME),
+        ))
+    iq = np.concatenate(chunks)
+    tx_rrc_time = time.perf_counter() - started
+    source_sha256 = (hashlib.sha256(SOURCE.read_bytes()).hexdigest()
+                     if pixels_override is None else hashlib.sha256(
+                         b"".join(item.tobytes() for item in pixels)
+                     ).hexdigest())
+    return VideoFixture(
+        iq, frames, tuple(intervals),
+        tuple(hashlib.sha256(frame.packet).hexdigest() for frame in phy_frames),
+        clean_hashes, source_sha256,
+        codec["settings"],
+        {"source_and_h264_encode": encode_time, "reference_h264_diagnostic": reference_decode_time,
+         "outer_fec_and_phy": outer_phy_time, "tx_rrc_and_iq": tx_rrc_time},
+    )
+
+
+def _loss_indices(fixture: VideoFixture, scenario: str, seed: int) -> tuple[set[int], int | None]:
+    random_generator = random.Random(seed)
+    if scenario == "clean":
+        return set(), None
+    if scenario in ("random_1", "random_5"):
+        probability = int(scenario.split("_")[1]) / 100
+        return {item.index for item in fixture.transmissions
+                if random_generator.random() < probability}, None
+    if scenario in ("keyframe_double", "keyframe_triple"):
+        count = 2 if scenario == "keyframe_double" else 3
+        by_block: dict[int, list[Transmission]] = {}
+        for item in fixture.transmissions:
+            if item.source_is_keyframe and item.source_frame_id is not None \
+                    and item.source_frame_id >= 5:
+                by_block.setdefault(item.block_id, []).append(item)
+        selected = next((items for items in by_block.values() if len(items) >= count), None)
+        if selected is None:
+            raise ValueError("не найден блок опорного кадра для выбранных потерь")
+        return {item.index for item in selected[:count]}, selected[0].source_frame_id
+    if scenario == "corrupt_packet":
+        selected = next(item for item in fixture.transmissions
+                        if item.source_frame_id is not None and item.source_frame_id >= 2)
+        return {selected.index}, selected.source_frame_id
+    raise ValueError("неизвестный сценарий стираний")
+
+
+def _packet_end_ms(packet_start_sample: float, full_packet_bytes: int) -> int:
+    body_symbols = math.ceil(punctured_fec.coded_bit_count(full_packet_bytes * 8, "2/3") / 2)
+    symbols = PREAMBLE_SYMBOLS + PHY_HEADER_SYMBOLS + body_symbols
+    # Начало в согласованном фильтре сдвинуто на длину двух RRC-фильтров минус один.
+    sample = packet_start_sample - (len(rrc_taps()) - 1) + symbols * 16
+    return math.ceil(sample * 1000 / SAMPLE_RATE_HZ)
+
+
+def assemble_video(fixture: VideoFixture, radio: IqStreamResult,
+                   lost_keyframe_id: int | None = None) -> dict:
+    """Подаёт только CRC-valid payload в FEC, сборщик и настоящий H.264-декодер."""
+    collector = VideoErasureCollector(STREAM_ID)
+    reassembler = FrameReassembler(STREAM_ID, 1012, max_age_ms=300)
+    decoder = av.CodecContext.create("h264", "r")
+    decoder.open()
+    displayed: list[tuple[int, int, int]] = []
+    completed_ids: list[int] = []
+    completed_data: list[tuple[int, bytes]] = []
+    decoded_frames: list[tuple[int, av.VideoFrame]] = []
+    decoder_errors = 0
+    erasure_time = 0.0
+    transport_time = 0.0
+    h264_time = 0.0
+    synchronized = False
+    last_displayed_id: int | None = None
+    seen_ids: set[int] = set()
+    for packet in radio.packets:
+        now_ms = _packet_end_ms(packet.start_sample, len(packet.packet))
+        started = time.perf_counter()
+        deliveries = collector.offer(packet.payload, now_ms)
+        erasure_time += time.perf_counter() - started
+        for delivery in deliveries:
+            age = now_ms - delivery.capture_ms
+            if not 0 <= age < 300:
+                raise AssertionError("видеофрагмент вышел за срок")
+            started = time.perf_counter()
+            complete = reassembler.offer_fragment(replace(delivery.fragment, age_ms=age), now_ms)
+            transport_time += time.perf_counter() - started
+            if complete is None:
+                continue
+            completed_ids.append(complete.frame_id)
+            completed_data.append((complete.frame_id, complete.data))
+            if complete.frame_type == FrameType.KEYFRAME:
+                decoder = av.CodecContext.create("h264", "r")
+                decoder.open()
+                synchronized = True
+            elif last_displayed_id is None or complete.frame_id != last_displayed_id + 1:
+                synchronized = False
+            if not synchronized:
+                continue
+            started = time.perf_counter()
+            encoded = av.Packet(complete.data)
+            encoded.pts = encoded.dts = complete.frame_id
+            encoded.time_base = Fraction(1, FPS)
+            try:
+                decoded = decoder.decode(encoded)
+            except FFmpegError:
+                h264_time += time.perf_counter() - started
+                decoder_errors += 1
+                synchronized = False
+                continue
+            if len(decoded) != 1 or decoded[0].pts != complete.frame_id:
+                h264_time += time.perf_counter() - started
+                decoder_errors += 1
+                synchronized = False
+                continue
+            h264_time += time.perf_counter() - started
+            decoded_frames.append((complete.frame_id, decoded[0]))
+            if complete.frame_id in seen_ids:
+                raise AssertionError("старый кадр повторно показан")
+            seen_ids.add(complete.frame_id)
+            displayed.append((complete.frame_id, now_ms, now_ms - complete.capture_ms))
+            last_displayed_id = complete.frame_id
+    collector.finish()
+    started = time.perf_counter()
+    wrong_encoded_frames = sum(
+        not 0 <= frame_id < len(fixture.frames)
+        or data != fixture.frames[frame_id].data
+        for frame_id, data in completed_data
+    )
+    wrong_decoded_frames = sum(
+        hashlib.sha256(frame.to_ndarray(format="rgb24").tobytes()).hexdigest()
+        != fixture.clean_pixel_sha256.get(frame_id)
+        for frame_id, frame in decoded_frames
+    )
+    diagnostic_time = time.perf_counter() - started
+    ages = [age for _, _, age in displayed]
+    recovery_ms = None
+    if lost_keyframe_id is not None:
+        if lost_keyframe_id in seen_ids:
+            recovery_ms = 0
+        else:
+            later = next((arrival - fixture.frames[lost_keyframe_id].capture_ms
+                          for frame_id, arrival, _ in displayed
+                          if frame_id > lost_keyframe_id
+                          and fixture.frames[frame_id].frame_type == FrameType.KEYFRAME), None)
+            recovery_ms = later
+    return {
+        "input_frames": len(fixture.frames),
+        "unique_decoded_frames": len(displayed),
+        "displayed_frame_ids": [frame_id for frame_id, _, _ in displayed],
+        "displayed_fps": len(displayed) * FPS / len(fixture.frames),
+        "frame_loss_rate": 1 - len(displayed) / len(fixture.frames),
+        "p50_frame_age_ms": float(np.percentile(ages, 50)) if ages else None,
+        "p95_frame_age_ms": float(np.percentile(ages, 95)) if ages else None,
+        "max_frame_age_ms": max(ages) if ages else None,
+        "stale_drops": collector.stale_rejections + reassembler.expired_frames,
+        "recovered_fragments": collector.parity_recoveries,
+        "unrecoverable_blocks": collector.unrecoverable_blocks,
+        "completed_frame_ids": completed_ids,
+        "wrong_encoded_frames": wrong_encoded_frames,
+        "wrong_decoded_frames": wrong_decoded_frames,
+        "h264_decoder_errors": decoder_errors,
+        "keyframe_recovery_ms": recovery_ms,
+        "old_frame_replay": displayed != sorted(displayed, key=lambda item: item[0]),
+        "erasure_fec_time_s": erasure_time,
+        "video_transport_time_s": transport_time,
+        "h264_decode_time_s": h264_time,
+        "diagnostic_verification_time_s": diagnostic_time,
+    }
+
+
+def run_case(
+    fixture: VideoFixture,
+    scenario: str = "clean",
+    *,
+    seed: int = 20260917,
+    ebn0_db: float = 9.0,
+    decoder_backend: str = "numba",
+) -> dict:
+    """Потери создаются в IQ; приёмник получает только повреждённый поток."""
+    selected, lost_keyframe_id = _loss_indices(fixture, scenario, seed)
+    started = time.perf_counter()
+    iq = impair(fixture.iq, cfo_hz=20, phase_degrees=15,
+                ebn0_db=ebn0_db, seed=seed)
+    if selected:
+        random_generator = np.random.default_rng(seed + 1)
+        for index in selected:
+            item = fixture.transmissions[index]
+            if scenario == "corrupt_packet":
+                sigma = 2.0
+                count = item.end_sample - item.start_sample
+                iq[item.start_sample:item.end_sample] = sigma * (
+                    random_generator.standard_normal(count)
+                    + 1j * random_generator.standard_normal(count)
+                )
+            else:
+                iq[item.start_sample:item.end_sample] = 0
+    channel_time = time.perf_counter() - started
+    started = time.perf_counter()
+    radio = receive_iq_stream(iq, decoder_backend=decoder_backend)
+    radio_time = time.perf_counter() - started
+    started = time.perf_counter()
+    video = assemble_video(fixture, radio, lost_keyframe_id)
+    video_time = time.perf_counter() - started
+    video_pipeline_time = video_time - video["diagnostic_verification_time_s"]
+    sequence_numbers = [parse_packet(packet.packet).sequence_number for packet in radio.packets]
+    # Сравнение с эталоном только после CRC и видео-декодирования.
+    exact_packets = len(set(sequence_numbers)) == len(sequence_numbers) and all(
+        number < len(fixture.phy_packet_sha256)
+        for number in sequence_numbers
+    ) and all(
+        hashlib.sha256(packet.packet).hexdigest() == fixture.phy_packet_sha256[number]
+        for packet, number in zip(radio.packets, sequence_numbers, strict=True)
+    )
+    return {
+        "scenario": scenario, "seed": seed, "ebn0_db": ebn0_db,
+        "iq_duration_s": len(iq) / SAMPLE_RATE_HZ,
+        "scheduled_phy_packets": len(fixture.transmissions),
+        "iq_erased_or_corrupted_packets": len(selected),
+        "crc_valid_radio_packets": len(radio.packets),
+        "radio_candidates": radio.candidate_count,
+        "exact_crc_packets": exact_packets,
+        "channel_time_s": channel_time, "radio_time_s": radio_time,
+        "video_transport_and_decode_time_s": video_pipeline_time,
+        "diagnostic_verification_time_s": video["diagnostic_verification_time_s"],
+        "radio_realtime_factor": radio_time / (len(iq) / SAMPLE_RATE_HZ),
+        "iq_to_decoded_video_realtime_factor": (
+            radio_time + video_pipeline_time
+        ) / (len(iq) / SAMPLE_RATE_HZ),
+        "end_to_end_realtime_factor": (
+            sum(duration for stage, duration in fixture.preparation_time_s.items()
+                if stage != "reference_h264_diagnostic")
+            + channel_time + radio_time + video_pipeline_time
+        ) / (len(iq) / SAMPLE_RATE_HZ),
+        "radio_stage_time_s": radio.stage_time_s,
+        "video": video,
+    }