Create-knowledge-find-decode/decode_content/readers/video_frames.py

"""Video frame extraction fallback for multimodal readers."""
from __future__ import annotations

import logging
import re
import subprocess
from pathlib import Path
from typing import Optional

import httpx

from core.models import Card

logger = logging.getLogger(__name__)

SCENE_THRESHOLD = 0.3
MIN_GAP_SECONDS = 1.5
FALLBACK_MIN = 4
FALLBACK_TARGET = 8
SCALE_VF = "scale='min(720,iw)':-2"
REFERER = {
    "douyin": "https://www.douyin.com/",
    "kuaishou": "https://www.kuaishou.com/",
    "bilibili": "https://www.bilibili.com/",
    "shipinhao": "https://channels.weixin.qq.com/",
    "xiaohongshu": "https://www.xiaohongshu.com/",
}


class VideoFramesError(RuntimeError):
    pass


def _ffmpeg_bin() -> str:
    try:
        import imageio_ffmpeg

        return imageio_ffmpeg.get_ffmpeg_exe()
    except Exception:
        return "ffmpeg"


def _download(url: str, platform: str, dst: Path, timeout: float) -> None:
    headers = {"User-Agent": "Mozilla/5.0", "Referer": REFERER.get(platform, "")}
    with httpx.stream("GET", url, headers=headers, timeout=timeout, follow_redirects=True) as response:
        response.raise_for_status()
        with open(dst, "wb") as file:
            for chunk in response.iter_bytes():
                file.write(chunk)


def _run(cmd: list[str]) -> str:
    proc = subprocess.run(cmd, capture_output=True, text=True)
    return proc.stderr or ""


def _probe_duration(ffmpeg: str, video: Path) -> float:
    err = _run([ffmpeg, "-hide_banner", "-i", str(video)])
    match = re.search(r"Duration:\s*(\d+):(\d+):(\d+\.?\d*)", err)
    if not match:
        return 0.0
    hours, minutes, seconds = match.groups()
    return int(hours) * 3600 + int(minutes) * 60 + float(seconds)


def _scene_frames(
    ffmpeg: str,
    video: Path,
    out: Path,
    threshold: float,
    min_gap: float,
    max_frames: int,
) -> list[tuple[Path, Optional[float]]]:
    pattern = str(out / "scene_%04d.jpg")
    err = _run(
        [
            ffmpeg,
            "-hide_banner",
            "-i",
            str(video),
            "-vf",
            f"select='gt(scene,{threshold})',showinfo,{SCALE_VF}",
            "-vsync",
            "vfr",
            "-q:v",
            "3",
            pattern,
        ]
    )
    times = [float(ts) for ts in re.findall(r"pts_time:(\d+\.?\d*)", err)]
    files = sorted(out.glob("scene_*.jpg"))
    paired = list(zip(files, times + [None] * (len(files) - len(times))))
    kept: list[tuple[Path, Optional[float]]] = []
    last = None
    for path, timestamp in paired:
        if timestamp is not None and last is not None and timestamp - last < min_gap:
            continue
        kept.append((path, timestamp))
        if timestamp is not None:
            last = timestamp
        if len(kept) >= max_frames:
            break
    return kept


def _uniform_frames(
    ffmpeg: str,
    video: Path,
    out: Path,
    duration: float,
    count: int,
) -> list[tuple[Path, Optional[float]]]:
    if duration <= 0:
        duration = float(count)
    result: list[tuple[Path, Optional[float]]] = []
    for idx in range(count):
        timestamp = duration * (idx + 0.5) / count
        path = out / f"uni_{idx:04d}.jpg"
        _run(
            [
                ffmpeg,
                "-hide_banner",
                "-ss",
                f"{timestamp:.2f}",
                "-i",
                str(video),
                "-frames:v",
                "1",
                "-vf",
                SCALE_VF,
                "-q:v",
                "3",
                str(path),
            ]
        )
        if path.exists():
            result.append((path, round(timestamp, 2)))
    return result


def extract_frames(
    video_src: str,
    *,
    out_dir: str | Path,
    url_prefix: str,
    platform: str = "",
    max_frames: int = 12,
    start_index: int = 1,
    timeout: float = 60.0,
) -> list[Card]:
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    ffmpeg = _ffmpeg_bin()

    src = Path(video_src)
    if not src.exists():
        src = out / "_src.mp4"
        _download(video_src, platform, src, timeout)

    frames = _scene_frames(ffmpeg, src, out, SCENE_THRESHOLD, MIN_GAP_SECONDS, max_frames)
    if len(frames) < FALLBACK_MIN:
        logger.info("scene frames only %d, falling back to uniform sampling", len(frames))
        duration = _probe_duration(ffmpeg, src)
        frames = _uniform_frames(ffmpeg, src, out, duration, FALLBACK_TARGET)
    frames = frames[:max_frames]
    if not frames:
        raise VideoFramesError(f"未能从视频抽出任何帧: {video_src}")

    return [
        Card(index=start_index + idx, kind="frame", url=f"{url_prefix.rstrip('/')}/{path.name}", timestamp=timestamp)
        for idx, (path, timestamp) in enumerate(frames)
    ]