Create-knowledge-find-decode/creation_knowledge/integrations/video_frames.py

"""视频抽帧：把视频转成一组"卡片"（帧），供多模态提取 + 溯源使用。

策略（见 技术文档/技术架构.md §12）：
  主策略 = 场景切换检测（ffmpeg select='gt(scene,T)'），知识类视频常是字幕卡/场景切换，
           正好对应"卡片"；附带 showinfo 解析每帧时间戳。
  护栏   = 最小间隔去近重复、最多 max_frames 控成本、下采样到长边≤720。
  回退   = 场景帧 < FALLBACK_MIN 时，按时长均匀采样 FALLBACK_TARGET 帧。

无 ffprobe 依赖：时长从 ffmpeg stderr 的 "Duration:" 解析。
帧文件写到 out_dir；Card.url = f"{url_prefix}/{文件名}"，由 API 的 /frames 静态服务。
"""
from __future__ import annotations

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

import httpx

from creation_knowledge.models import Card

logger = logging.getLogger(__name__)

SCENE_THRESHOLD = 0.3      # 场景切换阈值（电影 30° 规则）
MIN_GAP_SECONDS = 1.5      # 相邻保留帧最小间隔，去突发近重复
FALLBACK_MIN = 4           # 场景帧少于此数则触发均匀采样回退
FALLBACK_TARGET = 8        # 回退时采样的帧数
SCALE_VF = "scale='min(720,iw)':-2"  # 长边下采样到 ≤720

# 各平台下载视频需要的 Referer
_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 r:
        r.raise_for_status()
        with open(dst, "wb") as f:
            for chunk in r.iter_bytes():
                f.write(chunk)


def _run(cmd: list[str]) -> str:
    """跑 ffmpeg，返回 stderr（ffmpeg 把信息打到 stderr）。"""
    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)])
    m = re.search(r"Duration:\s*(\d+):(\d+):(\d+\.?\d*)", err)
    if not m:
        return 0.0
    h, mn, s = m.groups()
    return int(h) * 3600 + int(mn) * 60 + float(s)


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(t) for t 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, ts in paired:
        if ts is not None and last is not None and ts - last < min_gap:
            continue
        kept.append((path, ts))
        if ts is not None:
            last = ts
        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)  # 兜底，避免全 0
    result: list[tuple[Path, Optional[float]]] = []
    for k in range(count):
        ts = duration * (k + 0.5) / count
        path = out / f"uni_{k:04d}.jpg"
        _run([ffmpeg, "-hide_banner", "-ss", f"{ts:.2f}", "-i", str(video),
              "-frames:v", "1", "-vf", SCALE_VF, "-q:v", "3", str(path)])
        if path.exists():
            result.append((path, round(ts, 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]:
    """把视频抽成帧卡片。video_src 可为 URL 或本地文件路径。返回 Card(kind='frame')。"""
    out = Path(out_dir)
    out.mkdir(parents=True, exist_ok=True)
    ffmpeg = _ffmpeg_bin()

    src = Path(video_src)
    if not src.exists():  # 是 URL，先下载
        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("场景帧仅 %d，回退均匀采样", len(frames))
        dur = _probe_duration(ffmpeg, src)
        frames = _uniform_frames(ffmpeg, src, out, dur, FALLBACK_TARGET)
    frames = frames[:max_frames]
    if not frames:
        raise VideoFramesError(f"未能从视频抽出任何帧: {video_src}")

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