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
|