feat: 初始化 doco 子项目 P1 阶段(视频双路拆分预处理)

2026-06-12 16:31:21 +08:00
parent 70e2647fae
commit 0e6322ad97
10 changed files with 1149 additions and 13 deletions
@@ -0,0 +1,353 @@
+# -*- coding: utf-8 -*-
+"""
+视频双路拆分 - P1 核心模块
+=================================================
+功能:
+  A 路:视频帧 → pHash 变化检测 → OCR → B 稿 txt
+  B 路:视频 → 16kHz/单声道/16bit WAV
+
+不引入 ffmpeg-python 等 wrapper,只用 subprocess 调系统 ffmpeg。
+"""
+
+import hashlib
+import json
+import os
+import shutil
+import subprocess
+import tempfile
+from pathlib import Path
+from typing import Dict, List, Tuple, Optional
+
+from PIL import Image
+import imagehash
+
+
+# ========================================================================
+# 凭证(从环境变量读取,供 OCR 调用 DeepSeek Vision)
+# ========================================================================
+
+DEEPSEEK_API_KEY = os.environ.get("DEEPSEEK_API_KEY", "").strip()
+
+
+# ========================================================================
+# FFmpeg 封装
+# ========================================================================
+
+
+def check_ffmpeg():
+    """检查 ffmpeg 是否在 PATH 中"""
+    result = shutil.which("ffmpeg")
+    if result is None:
+        raise RuntimeError(
+            "ffmpeg 未找到,请先安装 ffmpeg 并加入 PATH。"
+            "下载地址: https://ffmpeg.org/download.html"
+        )
+    return result
+
+
+def extract_frames(
+    video_path: Path,
+    output_dir: Path,
+    fps: int = 1,
+) -> List[Tuple[int, int, Path]]:
+    """
+    按固定 fps 抽帧
+    返回: [(frame_index, timestamp_ms, image_path), ...]
+    """
+    check_ffmpeg()
+
+    frames_dir = output_dir / "frames"
+    frames_dir.mkdir(parents=True, exist_ok=True)
+
+    # ffmpeg 抽帧,格式 frame_%04d.png
+    frame_pattern = str(frames_dir / "frame_%04d.png")
+    cmd = [
+        "ffmpeg",
+        "-i", str(video_path),
+        "-vf", f"fps={fps}",
+        "-q:v", "2",  # JPEG 质量
+        frame_pattern,
+        "-y",  # 覆盖
+    ]
+
+    result = subprocess.run(
+        cmd,
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        raise RuntimeError(f"ffmpeg 抽帧失败: {result.stderr}")
+
+    # 收集抽出的帧
+    frames = []
+    for i, f in enumerate(sorted(frames_dir.glob("frame_*.png"))):
+        # 时间戳:第 i 帧就是 i 秒
+        timestamp_ms = i * 1000
+        frames.append((i, timestamp_ms, f))
+
+    return frames
+
+
+def extract_audio(
+    video_path: Path,
+    output_path: Path,
+) -> Path:
+    """
+    用 ffmpeg 提取音频,转为 16kHz/单声道/16bit WAV
+    """
+    check_ffmpeg()
+
+    cmd = [
+        "ffmpeg",
+        "-i", str(video_path),
+        "-ac", "1",           # 单声道
+        "-ar", "16000",       # 16kHz
+        "-sample_fmt", "s16", # 16bit
+        str(output_path),
+        "-y",
+    ]
+
+    result = subprocess.run(
+        cmd,
+        capture_output=True,
+        text=True,
+    )
+    if result.returncode != 0:
+        raise RuntimeError(f"ffmpeg 音频提取失败: {result.stderr}")
+
+    return output_path
+
+
+# ========================================================================
+# pHash 变化检测
+# ========================================================================
+
+
+def compute_phash(image_path: Path) -> str:
+    """计算图片的 pHash,返回 hex 字符串"""
+    img = Image.open(image_path)
+    ph = imagehash.phash(img)
+    return str(ph)
+
+
+def find_keyframes(
+    frames: List[Tuple[int, int, Path]],
+    threshold: int = 8,
+) -> List[Dict]:
+    """
+    基于 pHash 海明距离找出字幕变化的关键帧
+
+    算法:
+    - 第一帧总是关键帧
+    - 后续帧:如果与上一个关键帧的 pHash 海明距离 > threshold,则是新关键帧
+
+    threshold: 海明距离阈值,默认 8
+    """
+    if not frames:
+        return []
+
+    keyframes = []
+    last_keyframe_phash = None
+
+    for frame_index, timestamp_ms, image_path in frames:
+        phash = compute_phash(image_path)
+
+        is_keyframe = False
+        if last_keyframe_phash is None:
+            # 第一帧总是关键帧
+            is_keyframe = True
+        else:
+            # 计算海明距离
+            hamming = hamming_distance(last_keyframe_phash, phash)
+            if hamming > threshold:
+                is_keyframe = True
+
+        if is_keyframe:
+            keyframes.append({
+                "frame_index": frame_index,
+                "timestamp_ms": timestamp_ms,
+                "frame_image_path": str(image_path),
+                "phash": phash,
+                "ocr_text": "",  # P2 调用 DeepSeek Vision 填充
+            })
+            last_keyframe_phash = phash
+
+    return keyframes
+
+
+def hamming_distance(s1: str, s2: str) -> int:
+    """计算两个 hex pHash 字符串的海明距离"""
+    if len(s1) != len(s2):
+        # pHash 长度不一致,取较长字符串的长度作为海明距离上限
+        return max(len(s1), len(s2))
+    return sum(c1 != c2 for c1, c2 in zip(s1, s2))
+
+
+# ========================================================================
+# OCR 接口(P2 实现,目前返回占位)
+# ========================================================================
+
+
+def ocr_frame(image_path: Path) -> str:
+    """
+    识别帧内文字,返回纯文本
+
+    P1: 返回占位文本
+    P2: 调用 DeepSeek Vision API
+    """
+    if not DEEPSEEK_API_KEY:
+        # 无 API Key,返回占位
+        return f"[OCR待填充 frame={image_path.name}]"
+
+    # P2 实现:调用 DeepSeek Vision
+    # TODO: P2 实现
+    return f"[OCR待填充 frame={image_path.name}]"
+
+
+def ocr_keyframes(keyframes: List[Dict]) -> List[Dict]:
+    """对关键帧列表逐一调用 OCR"""
+    result = []
+    for kf in keyframes:
+        image_path = Path(kf["frame_image_path"])
+        ocr_text = ocr_frame(image_path)
+        kf_copy = kf.copy()
+        kf_copy["ocr_text"] = ocr_text
+        result.append(kf_copy)
+    return result
+
+
+# ========================================================================
+# B 稿格式化
+# ========================================================================
+
+
+def format_timestamp(ms: int) -> str:
+    """毫秒转 [Nm Ns] 格式"""
+    total_sec = ms // 1000
+    return f"{total_sec // 60}m{total_sec % 60}s"
+
+
+def build_b_manuscript(keyframes: List[Dict]) -> List[str]:
+    """
+    将关键帧 OCR 结果合并为 B 稿
+    合并相邻同文本的关键帧
+    """
+    lines = []
+    last_text = None
+
+    for kf in keyframes:
+        text = kf["ocr_text"].strip()
+        if not text:
+            continue
+
+        # 跳过占位文本
+        if text.startswith("[OCR待填充"):
+            text = ""
+
+        if text and text != last_text:
+            ts = format_timestamp(kf["timestamp_ms"])
+            lines.append(f"[{ts}] {text}")
+            last_text = text
+
+    return lines
+
+
+def write_b_manuscript(lines: List[str], output_path: Path) -> Path:
+    """写入 B 稿 txt"""
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    with open(output_path, "w", encoding="utf-8") as f:
+        f.write("\n".join(lines))
+    return output_path
+
+
+# ========================================================================
+# 主流程
+# ========================================================================
+
+
+def split_video(
+    video_path: Path,
+    output_dir: Path,
+    episode_id: str,
+    phash_threshold: int = 8,
+    fps: int = 1,
+) -> Dict[str, any]:
+    """
+    视频双路拆分主流程
+
+    参数:
+        video_path: 输入视频路径
+        output_dir: 输出目录(work/ 路径)
+        episode_id: 节目 ID
+        phash_threshold: pHash 海明距离阈值,默认 8
+        fps: 抽帧帧率,默认 1(每秒一帧)
+
+    返回:
+        {
+            "b_manuscript_path": Path,
+            "audio_path": Path,
+            "keyframes_path": Path,
+            "keyframe_count": int,
+        }
+    """
+    video_path = Path(video_path)
+    output_dir = Path(output_dir)
+
+    if not video_path.exists():
+        raise FileNotFoundError(f"视频文件不存在: {video_path}")
+
+    # 创建输出目录
+    output_dir.mkdir(parents=True, exist_ok=True)
+    frames_dir = output_dir / "frames"
+    frames_dir.mkdir(parents=True, exist_ok=True)
+
+    print(f"[video_split] 开始处理: {video_path.name}")
+    print(f"[video_split] 抽帧 fps={fps}, pHash threshold={phash_threshold}")
+
+    # ---- A 路:抽帧 + pHash 检测 + OCR ----
+    print("[video_split] A路:抽帧...")
+    frames = extract_frames(video_path, output_dir, fps=fps)
+    print(f"[video_split] 抽帧完成,共 {len(frames)} 帧")
+
+    print("[video_split] pHash 变化检测...")
+    keyframes = find_keyframes(frames, threshold=phash_threshold)
+    print(f"[video_split] 检测到 {len(keyframes)} 个关键帧")
+
+    print("[video_split] OCR 关键帧...")
+    keyframes = ocr_keyframes(keyframes)
+    print(f"[video_split] OCR 完成")
+
+    # ---- B 路:音频提取 ----
+    print("[video_split] B路:提取音频...")
+    audio_path = output_dir / "audio_16k.wav"
+    extract_audio(video_path, audio_path)
+    print(f"[video_split] 音频提取完成: {audio_path}")
+
+    # ---- 输出产物 ----
+    # B 稿
+    b_lines = build_b_manuscript(keyframes)
+    b_manuscript_path = output_dir / "b_manuscript.txt"
+    write_b_manuscript(b_lines, b_manuscript_path)
+    print(f"[video_split] B稿写入: {b_manuscript_path} ({len(b_lines)} 行)")
+
+    # 关键帧索引 JSON
+    keyframes_data = {
+        "video_path": str(video_path),
+        "fps_sampled": fps,
+        "phash_threshold": phash_threshold,
+        "keyframes": keyframes,
+    }
+    keyframes_path = output_dir / "keyframes.json"
+    with open(keyframes_path, "w", encoding="utf-8") as f:
+        json.dump(keyframes_data, f, ensure_ascii=False, indent=2)
+    print(f"[video_split] 关键帧索引写入: {keyframes_path}")
+
+    # 清理临时帧文件(可选,保留供调试)
+    # shutil.rmtree(frames_dir)
+
+    return {
+        "b_manuscript_path": str(b_manuscript_path),
+        "audio_path": str(audio_path),
+        "keyframes_path": str(keyframes_path),
+        "keyframe_count": len(keyframes),
+    }