tmp/video-velocity-model/data_processing/tools/cal_fvd.py

"""Compute FVD between real .avi videos and generated mp4s.

Reads a manifest produced by `data_processing/sample_for_fvd.py` and writes a
`fvd_score.json` next to it. I3D feature extraction follows Latte's
`Latte/tools/metrics/frechet_video_distance.py` contract:
  - TorchScript I3D pretrained on Kinetics-400 (Dropbox URL, cached locally)
  - `rescale=True, resize=True, return_features=True`
  - input shape `[B, C=3, T, H, W]` float, value range ~ [0, 255]
"""

from __future__ import annotations

import argparse
import json
import sys
import urllib.request
from pathlib import Path
from typing import Iterable, List, Tuple

import numpy as np
import scipy.linalg
import torch
from tqdm import tqdm

REPO_ROOT = Path(__file__).resolve().parents[1]
if str(REPO_ROOT / "data_processing") not in sys.path:
    sys.path.append(str(REPO_ROOT / "data_processing"))


def _load_video_frames(video_path: Path, image_size: int):
    """Lazy import of forward.load_video_frames so unit tests don't pull decord / stage1."""
    try:
        from .forward import load_video_frames  # type: ignore
    except ImportError:
        from forward import load_video_frames  # type: ignore
    return load_video_frames(video_path, image_size)


I3D_URL = "https://www.dropbox.com/s/ge9e5ujwgetktms/i3d_torchscript.pt?dl=1"
I3D_KWARGS = dict(rescale=True, resize=True, return_features=True)
I3D_FEATURE_DIM = 400


def ensure_i3d_ckpt(path: Path) -> Path:
    path = path.expanduser()
    if path.exists():
        return path
    print(f"[fvd] downloading I3D torchscript to {path} ...", flush=True)
    path.parent.mkdir(parents=True, exist_ok=True)
    try:
        urllib.request.urlretrieve(I3D_URL, str(path))
    except Exception as exc:
        raise RuntimeError(
            f"Failed to download I3D ckpt from {I3D_URL}. Place the file at {path} manually."
        ) from exc
    return path


def load_i3d(path: Path, device: torch.device):
    model = torch.jit.load(str(path), map_location=device)
    model.eval()
    return model


def load_clip(video_path: Path, num_frames: int, image_size: int) -> torch.Tensor:
    """Read mp4/avi, take first `num_frames`, center-crop to image_size, return [3, T, H, W] in [0, 1]."""
    frames, _fps = _load_video_frames(video_path, image_size)  # [F, H, W, 3] uint8
    if frames.shape[0] < num_frames:
        raise ValueError(f"{video_path} has {frames.shape[0]} frames, needs >= {num_frames}")
    clip = frames[:num_frames]
    tensor = torch.from_numpy(clip).float() / 255.0  # [T, H, W, 3] in [0, 1]
    tensor = tensor.permute(3, 0, 1, 2).contiguous()  # [3, T, H, W]
    return tensor


class StreamingStats:
    """Streaming mean / covariance for [N, D] features. All math in float64."""

    def __init__(self, dim: int):
        self.n = 0
        self.sum = np.zeros(dim, dtype=np.float64)
        self.sumsq = np.zeros((dim, dim), dtype=np.float64)

    def update(self, x: np.ndarray) -> None:
        x = np.asarray(x, dtype=np.float64)
        if x.ndim != 2 or x.shape[1] != self.sum.shape[0]:
            raise ValueError(f"expected [B, {self.sum.shape[0]}], got {x.shape}")
        self.n += x.shape[0]
        self.sum += x.sum(0)
        self.sumsq += x.T @ x

    def finalize(self) -> Tuple[np.ndarray, np.ndarray]:
        if self.n == 0:
            raise RuntimeError("No samples accumulated.")
        mu = self.sum / self.n
        cov = self.sumsq / self.n - np.outer(mu, mu)
        return mu, cov


@torch.no_grad()
def collect_features(i3d, clips: torch.Tensor, device: torch.device) -> np.ndarray:
    """`clips`: [B, 3, T, H, W] in [0, 1]. Returns [B, 400] float64."""
    x = (clips.to(device, non_blocking=True) * 255.0).float()
    feats = i3d(x, **I3D_KWARGS)
    return feats.cpu().numpy().astype(np.float64)


def fvd_score(mu_r: np.ndarray, cov_r: np.ndarray, mu_g: np.ndarray, cov_g: np.ndarray) -> float:
    m = float(np.square(mu_g - mu_r).sum())
    prod = cov_g @ cov_r
    s, _ = scipy.linalg.sqrtm(prod, disp=False)
    if not np.isfinite(s).all():
        eps = 1e-6 * np.eye(cov_r.shape[0])
        s, _ = scipy.linalg.sqrtm((cov_g + eps) @ (cov_r + eps), disp=False)
    return float(np.real(m + np.trace(cov_g + cov_r - 2.0 * s)))


def batched_pairs(entries: List[dict], batch_size: int) -> Iterable[List[dict]]:
    for start in range(0, len(entries), batch_size):
        yield entries[start : start + batch_size]


def parse_args() -> argparse.Namespace:
    p = argparse.ArgumentParser(description="Compute FVD between real .avi videos and generated mp4s.")
    p.add_argument("--manifest", type=Path, required=True)
    p.add_argument("--output", type=Path, default=None, help="Defaults to manifest.parent / 'fvd_score.json'.")
    p.add_argument(
        "--i3d-ckpt",
        type=Path,
        default=Path.home() / ".cache" / "zst-project" / "i3d_torchscript.pt",
    )
    p.add_argument("--num-frames", type=int, default=16)
    p.add_argument("--image-size", type=int, default=256)
    p.add_argument("--batch-size", type=int, default=16)
    p.add_argument("--device", type=str, default="cuda:0")
    return p.parse_args()


def main() -> None:
    args = parse_args()
    manifest_path: Path = args.manifest
    entries: List[dict] = json.loads(manifest_path.read_text())
    if not entries:
        raise RuntimeError(f"Manifest {manifest_path} is empty.")

    device = torch.device(args.device)
    i3d = load_i3d(ensure_i3d_ckpt(args.i3d_ckpt), device)

    stats_real = StreamingStats(I3D_FEATURE_DIM)
    stats_gen = StreamingStats(I3D_FEATURE_DIM)

    skipped: List[dict] = []
    pbar = tqdm(total=len(entries), desc="fvd")
    for batch in batched_pairs(entries, args.batch_size):
        real_clips = []
        gen_clips = []
        kept_batch = []
        for e in batch:
            try:
                real = load_clip(Path(e["avi_path"]), args.num_frames, args.image_size)
                gen = load_clip(Path(e["gen_path"]), args.num_frames, args.image_size)
            except (ValueError, OSError, RuntimeError) as exc:
                skipped.append({"entry": e, "error": str(exc)})
                pbar.update(1)
                continue
            real_clips.append(real)
            gen_clips.append(gen)
            kept_batch.append(e)

        if not kept_batch:
            continue

        real_tensor = torch.stack(real_clips, dim=0)  # [B, 3, T, H, W]
        gen_tensor = torch.stack(gen_clips, dim=0)
        stats_real.update(collect_features(i3d, real_tensor, device))
        stats_gen.update(collect_features(i3d, gen_tensor, device))
        pbar.update(len(kept_batch))
    pbar.close()

    if stats_real.n == 0:
        raise RuntimeError("No valid (real, gen) pairs processed.")
    if stats_real.n != stats_gen.n:
        raise RuntimeError(f"Real/gen sample counts disagree: {stats_real.n} vs {stats_gen.n}")

    mu_r, cov_r = stats_real.finalize()
    mu_g, cov_g = stats_gen.finalize()
    fvd = fvd_score(mu_r, cov_r, mu_g, cov_g)

    out_path = args.output or (manifest_path.parent / "fvd_score.json")
    out_path.write_text(
        json.dumps(
            {
                "fvd": fvd,
                "num_videos": stats_real.n,
                "num_skipped": len(skipped),
                "num_frames": args.num_frames,
                "image_size": args.image_size,
                "i3d_ckpt": str(args.i3d_ckpt),
                "manifest": str(manifest_path),
                "skipped_examples": skipped[:8],
            },
            indent=2,
        )
    )
    print(f"FVD = {fvd:.4f} on {stats_real.n} videos (skipped {len(skipped)}) -> {out_path}", flush=True)


if __name__ == "__main__":
    main()