simonkoson
354 lines
9.8 KiB
Python
354 lines
9.8 KiB
Python
# -*- 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),
|
|
}
|