101 lines
3.9 KiB
Python
101 lines
3.9 KiB
Python
|
"""音视频转写：faster-whisper 本地模型 + ffmpeg 抽音轨。
|
|||
|
|
输出与 .srt 同格式的「去时间戳、按语义间隙合并」文本，直接进入既有 RAG 分块管线。
|
|||
|
|
"""
|
|||
|
|
from __future__ import annotations
|
|||
|
|
|
|||
|
|
import os
|
|||
|
|
import re
|
|||
|
|
import subprocess
|
|||
|
|
import threading
|
|||
|
|
from pathlib import Path
|
|||
|
|
from .parse import clean_text
|
|||
|
|
|
|||
|
|
MEDIA_SUFFIXES = {".mp4", ".mov", ".m4a", ".mp3", ".wav", ".flv", ".mkv", ".avi", ".webm", ".aac", ".ogg", ".flac"}
|
|||
|
|
|
|||
|
|
# 模型目录：优先环境变量，其次 modelscope 缓存里的 gpustack/faster-whisper-small
|
|||
|
|
_MODEL_DIR_CANDIDATES = [
|
|||
|
|
os.environ.get("WHISPER_MODEL_DIR", ""),
|
|||
|
|
str(Path.home() / ".cache/modelscope/models/gpustack--faster-whisper-small/snapshots/master"),
|
|||
|
|
]
|
|||
|
|
_model = None
|
|||
|
|
_model_lock = threading.Lock()
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _get_model():
|
|||
|
|
global _model
|
|||
|
|
if _model is not None:
|
|||
|
|
return _model
|
|||
|
|
with _model_lock:
|
|||
|
|
if _model is not None:
|
|||
|
|
return _model
|
|||
|
|
from faster_whisper import WhisperModel
|
|||
|
|
for cand in _MODEL_DIR_CANDIDATES:
|
|||
|
|
if cand and Path(cand).exists():
|
|||
|
|
_model = WhisperModel(cand, device="cpu", compute_type="int8", local_files_only=True)
|
|||
|
|
return _model
|
|||
|
|
_model = WhisperModel("small", device="cpu", compute_type="int8")
|
|||
|
|
return _model
|
|||
|
|
|
|||
|
|
|
|||
|
|
def is_media(path: Path) -> bool:
|
|||
|
|
return path.suffix.lower() in MEDIA_SUFFIXES
|
|||
|
|
|
|||
|
|
|
|||
|
|
def _probe_duration(path: Path) -> float:
|
|||
|
|
try:
|
|||
|
|
r = subprocess.run(
|
|||
|
|
["ffprobe", "-v", "quiet", "-show_entries", "format=duration",
|
|||
|
|
"-of", "csv=p=0", str(path)],
|
|||
|
|
capture_output=True, text=True, timeout=30)
|
|||
|
|
return float(r.stdout.strip() or 0)
|
|||
|
|
except Exception:
|
|||
|
|
return 0.0
|
|||
|
|
|
|||
|
|
|
|||
|
|
def transcribe_media(path: Path, progress_cb=None) -> list[dict]:
|
|||
|
|
"""音视频 → 语义段落列表 [{text, meta:{heading}}]。
|
|||
|
|
超长音先压缩码率降体积；VAD 过滤静音；按标点+间隙合并成语义段。"""
|
|||
|
|
src = Path(path)
|
|||
|
|
work = src
|
|||
|
|
# 大文件先转 16k mono wav（whisper 输入要求，同时把视频降为纯音轨）
|
|||
|
|
if src.suffix.lower() not in (".wav",) or _probe_duration(src) > 0:
|
|||
|
|
work = src.with_suffix(".wav.wav")
|
|||
|
|
cmd = ["ffmpeg", "-y", "-i", str(src), "-vn", "-ar", "16000", "-ac", "1",
|
|||
|
|
"-b:a", "64k", str(work)]
|
|||
|
|
r = subprocess.run(cmd, capture_output=True, text=True, timeout=3600)
|
|||
|
|
if r.returncode != 0 or not work.exists() or work.stat().st_size == 0:
|
|||
|
|
raise ValueError(f"ffmpeg 抽音轨失败: {r.stderr[-200:]}")
|
|||
|
|
try:
|
|||
|
|
model = _get_model()
|
|||
|
|
segments, _info = model.transcribe(str(work), language="zh", vad_filter=True,
|
|||
|
|
vad_parameters={"min_silence_duration_ms": 700})
|
|||
|
|
lines: list[str] = []
|
|||
|
|
for seg in segments:
|
|||
|
|
t = (seg.text or "").strip()
|
|||
|
|
if t:
|
|||
|
|
lines.append(t)
|
|||
|
|
if progress_cb and len(lines) % 20 == 0:
|
|||
|
|
progress_cb(len(lines))
|
|||
|
|
if not lines:
|
|||
|
|
return []
|
|||
|
|
# 合并成语义段：约 15 句或 1200 字一段（与既有 2200 字分块目标匹配）
|
|||
|
|
merged, buf, size = [], [], 0
|
|||
|
|
for t in lines:
|
|||
|
|
buf.append(t)
|
|||
|
|
size += len(t)
|
|||
|
|
if size >= 1200 or t.endswith(("。", "！", "？", "!", "?")) and size >= 600:
|
|||
|
|
merged.append(" ".join(buf))
|
|||
|
|
buf, size = [], 0
|
|||
|
|
if buf:
|
|||
|
|
merged.append(" ".join(buf))
|
|||
|
|
# 转写结果也经过同一套结构清洗：去掉字幕式序号/时间戳、不可见字符和句首独立口头填充词，
|
|||
|
|
# 保留正文中的日期、金额、型号等数字。
|
|||
|
|
return [{"text": clean_text(m), "meta": {"heading": src.stem}}
|
|||
|
|
for m in merged if clean_text(m)]
|
|||
|
|
finally:
|
|||
|
|
if work != src:
|
|||
|
|
try:
|
|||
|
|
work.unlink(missing_ok=True)
|
|||
|
|
except OSError:
|
|||
|
|
pass
|