remove-ai-watermarks/scripts/video_fidelity_probe.py

"""End-to-end fidelity of a DELIVERED video against its source.

The engine's own ``psnr_db`` is measured against the already-resized frame and
before the encoder, so it reports the VAE round trip plus latent noise and
nothing else. The downscale, the frame decimation and the H.264 encode -- the
three steps that actually cost the user picture -- are invisible to it, and no
metric computed inside the streaming loop can see them, because the candidate
frame is captured before it is written to the encoder pipe.

This script measures what the engine cannot: it decodes the delivered file after
muxing, upscales each frame back to the source geometry, and scores it against
the untouched source frame it came from. That reference is deliberately harsh -
detail the downscale destroyed is unrecoverable, so the absolute number is
dominated by content the pipeline never had a chance to keep. Use it to RANK
configurations against a shared reference, not to attribute loss to one stage.

It streams, for the same reason the engine does: holding a 1080p clip plus its
upscaled candidate and flow maps in memory runs to gigabytes and grows with clip
length. Peak here is a handful of frames regardless of duration.

It is a research tool, not part of the shipped path, and it measures fidelity
only. Nothing here is a watermark verdict: only a provider-oracle row is.

    uv run python scripts/video_fidelity_probe.py source.mp4 out-512.mp4 out-768.mp4
"""

from __future__ import annotations

import json
import logging
import math
import sys
from pathlib import Path
from typing import TYPE_CHECKING, Any

import click
import cv2
import numpy as np

from remove_ai_watermarks.video_invisible import _iter_sampled_frames, _probe_video
from remove_ai_watermarks.video_temporal import _backward_map, _motion_residual

sys.path.insert(0, str(Path(__file__).parent))

from invisible_quality_audit import _ssim  # reuse, do not reimplement a fourth SSIM

if TYPE_CHECKING:
    from collections.abc import Iterator

    from numpy.typing import NDArray

log = logging.getLogger(__name__)


def _delivered_geometry(path: Path) -> tuple[int, int, float]:
    capture = cv2.VideoCapture(str(path))
    if not capture.isOpened():
        raise click.ClickException(f"Could not open video: {path}")
    try:
        width = round(capture.get(cv2.CAP_PROP_FRAME_WIDTH))
        height = round(capture.get(cv2.CAP_PROP_FRAME_HEIGHT))
        fps = float(capture.get(cv2.CAP_PROP_FPS))
    finally:
        capture.release()
    if width <= 0 or height <= 0 or fps <= 0.0:
        raise click.ClickException(f"Video has no usable geometry or frame rate: {path}")
    return width, height, fps


def _iter_frames(path: Path) -> Iterator[NDArray[Any]]:
    capture = cv2.VideoCapture(str(path))
    if not capture.isOpened():
        raise click.ClickException(f"Could not open video: {path}")
    try:
        while True:
            ok, frame = capture.read()
            if not ok:
                return
            yield frame
    finally:
        capture.release()


def _measure(
    source: Path,
    delivered: Path,
    *,
    source_geometry: tuple[int, int, float],
    duration: float | None,
) -> dict[str, Any]:
    source_width, source_height, source_fps = source_geometry
    width, height, fps = _delivered_geometry(delivered)

    # Drive the source through the engine's own sampler at the source's geometry,
    # so the frame the probe compares against is the frame the engine regenerated.
    # Importing the rule is what keeps the pairing correct: a count check cannot
    # catch a selection rule that reorders frames without changing how many.
    reference_frames = _iter_sampled_frames(
        source,
        source_fps=source_fps,
        duration=duration,
        effective_fps=fps,
        size=(source_width, source_height),
    )

    squared_error = 0.0
    pixel_count = 0
    ssim_scores: list[float] = []
    temporal_baseline = 0.0
    temporal_candidate = 0.0
    previous_gray: NDArray[Any] | None = None
    previous_reference_f32: NDArray[Any] | None = None
    previous_candidate_f32: NDArray[Any] | None = None
    frame_count = 0
    needs_upscale = (width, height) != (source_width, source_height)

    reference_iter = iter(reference_frames)
    delivered_iter = _iter_frames(delivered)
    while True:
        reference = next(reference_iter, None)
        delivered_frame = next(delivered_iter, None)
        if reference is None and delivered_frame is None:
            break
        if reference is None or delivered_frame is None:
            raise click.ClickException(
                f"{delivered.name} and the sampled source ran out at different points after "
                f"{frame_count} frames. Pass --duration to match the prefix this output was "
                f"produced from, or check that it came from {source.name}."
            )
        candidate = (
            cv2.resize(delivered_frame, (source_width, source_height), interpolation=cv2.INTER_LANCZOS4)
            if needs_upscale
            else delivered_frame
        )
        reference_f32 = reference.astype(np.float32)
        candidate_f32 = candidate.astype(np.float32)
        difference = reference_f32 - candidate_f32
        squared_error += float(np.sum(difference * difference, dtype=np.float64))
        pixel_count += reference.size
        ssim_scores.append(
            _ssim(cv2.cvtColor(reference, cv2.COLOR_BGR2GRAY), cv2.cvtColor(candidate, cv2.COLOR_BGR2GRAY))
        )

        current_gray = cv2.cvtColor(reference, cv2.COLOR_BGR2GRAY)
        if previous_gray is not None and previous_reference_f32 is not None and previous_candidate_f32 is not None:
            frame_maps = _backward_map(current_gray, previous_gray)
            temporal_baseline += _motion_residual(reference_f32, previous_reference_f32, frame_maps)
            temporal_candidate += _motion_residual(candidate_f32, previous_candidate_f32, frame_maps)
        previous_gray = current_gray
        previous_reference_f32 = reference_f32
        previous_candidate_f32 = candidate_f32
        frame_count += 1

    if frame_count < 2:
        raise click.ClickException(f"{delivered.name} paired fewer than two frames against {source.name}")

    mse = squared_error / pixel_count
    size_bytes = delivered.stat().st_size
    return {
        "file": delivered.name,
        "source": source.name,
        "source_width": source_width,
        "source_height": source_height,
        "source_fps": round(source_fps, 4),
        "width": width,
        "height": height,
        "fps": round(fps, 4),
        "frames": frame_count,
        "pixel_ratio": round((width * height) / (source_width * source_height), 4),
        "source_psnr_db": math.inf if mse == 0.0 else round(20.0 * math.log10(255.0 / math.sqrt(mse)), 4),
        "source_ssim": round(float(np.mean(ssim_scores)), 4),
        "temporal_residual_ratio": round(temporal_candidate / max(temporal_baseline, 1e-6), 4),
        "size_bytes": size_bytes,
        # The mux copies the source audio track verbatim, so this is a container
        # bitrate. It still ranks candidates at fixed crf; it is not a video bitrate.
        "file_bitrate_kbps": round(size_bytes * 8.0 / (frame_count / fps) / 1000.0, 1),
    }


@click.command()
@click.argument("source", type=click.Path(exists=True, dir_okay=False, path_type=Path))
@click.argument("delivered", nargs=-1, required=True, type=click.Path(exists=True, dir_okay=False, path_type=Path))
@click.option(
    "--duration",
    type=click.FloatRange(min=0.1),
    default=None,
    help="Score only the first N seconds of the source, matching a trimmed sweep candidate.",
)
@click.option("--json-out", type=click.Path(dir_okay=False, path_type=Path), help="Also write the rows as JSON.")
def main(source: Path, delivered: tuple[Path, ...], duration: float | None, json_out: Path | None) -> None:
    """Score each DELIVERED file against SOURCE end to end."""
    logging.basicConfig(level=logging.INFO, format="%(message)s")
    source_geometry = _probe_video(source)
    log.info("Source %s: %dx%d at %.4f fps", source.name, *source_geometry)

    rows = [_measure(source, path, source_geometry=source_geometry, duration=duration) for path in delivered]
    for row in rows:
        log.info(
            "%s: %dx%d at %s fps, %.1f%% of source pixels, PSNR %s dB, SSIM %s, temporal %s, %s kbps",
            row["file"],
            row["width"],
            row["height"],
            row["fps"],
            row["pixel_ratio"] * 100.0,
            row["source_psnr_db"],
            row["source_ssim"],
            row["temporal_residual_ratio"],
            row["file_bitrate_kbps"],
        )
    if json_out is not None:
        json_out.write_text(json.dumps(rows, indent=2), encoding="utf-8")
        log.info("Wrote %s", json_out)


if __name__ == "__main__":
    main()