"""Transcribe the first audio stream of one local video with faster-whisper."""
|
|
from __future__ import annotations
|
|
import argparse
|
import json
|
import math
|
import os
|
import shutil
|
import subprocess
|
import sys
|
import tempfile
|
from dataclasses import dataclass
|
from pathlib import Path
|
from typing import Any, Callable, Iterable, Sequence
|
|
|
MODEL_NAME = "large-v3"
|
DEVICE = "cuda"
|
COMPUTE_TYPE = "float16"
|
SHORT_AUDIO_LIMIT_SECONDS = 1200.0
|
CHUNK_CORE_SECONDS = 1200.0
|
CHUNK_OVERLAP_SECONDS = 5.0
|
|
|
class MediaTranscriptionError(RuntimeError):
|
"""A user-facing error that should stop the command without a traceback."""
|
|
|
@dataclass(frozen=True)
|
class OutputPaths:
|
directory: Path
|
audio: Path
|
text: Path
|
srt: Path
|
json: Path
|
|
@property
|
def files(self) -> tuple[Path, Path, Path, Path]:
|
return (self.audio, self.text, self.srt, self.json)
|
|
|
@dataclass(frozen=True)
|
class RuntimeInfo:
|
cuda_device_count: int
|
cuda_compute_types: tuple[str, ...]
|
|
|
@dataclass(frozen=True)
|
class TranscriptSegment:
|
start: float
|
end: float
|
text: str
|
|
|
@dataclass(frozen=True)
|
class TranscriptResult:
|
language: str | None
|
language_probability: float | None
|
segments: tuple[TranscriptSegment, ...]
|
|
|
@dataclass(frozen=True)
|
class AudioChunk:
|
index: int
|
total: int
|
core_start: float
|
core_end: float
|
audio_start: float
|
audio_end: float
|
|
@property
|
def audio_duration(self) -> float:
|
return self.audio_end - self.audio_start
|
|
|
ProcessRunner = Callable[[Sequence[str]], subprocess.CompletedProcess[str]]
|
ModelLoader = Callable[[], tuple[Any, RuntimeInfo]]
|
|
|
def _run_process(command: Sequence[str]) -> subprocess.CompletedProcess[str]:
|
return subprocess.run(
|
list(command),
|
check=False,
|
capture_output=True,
|
text=True,
|
encoding="utf-8",
|
errors="replace",
|
)
|
|
|
def _error_detail(result: subprocess.CompletedProcess[str]) -> str:
|
detail = (result.stderr or result.stdout or "未返回错误详情").strip()
|
return detail[-1200:]
|
|
|
def _require_executable(name: str) -> str:
|
executable = shutil.which(name)
|
if executable is None:
|
raise MediaTranscriptionError(
|
f"未找到 {name}，请先安装并确保它位于当前 PATH。"
|
)
|
return executable
|
|
|
def _resolve_video(video_path: str | Path) -> Path:
|
video = Path(video_path).expanduser().resolve()
|
if not video.exists():
|
raise MediaTranscriptionError(f"视频文件不存在：{video}")
|
if not video.is_file():
|
raise MediaTranscriptionError(f"输入路径不是文件：{video}")
|
return video
|
|
|
def _build_output_paths(video: Path, output_dir: str | Path | None) -> OutputPaths:
|
if output_dir is None:
|
directory = video.with_name(f"{video.stem}.transcript")
|
else:
|
directory = Path(output_dir).expanduser().resolve()
|
|
stem = video.stem
|
return OutputPaths(
|
directory=directory,
|
audio=directory / f"{stem}.audio.flac",
|
text=directory / f"{stem}.txt",
|
srt=directory / f"{stem}.srt",
|
json=directory / f"{stem}.json",
|
)
|
|
|
def _ensure_output_available(outputs: OutputPaths) -> None:
|
if outputs.directory.exists() and not outputs.directory.is_dir():
|
raise MediaTranscriptionError(
|
f"输出路径已存在且不是目录：{outputs.directory}"
|
)
|
|
conflicts = [path for path in outputs.files if path.exists()]
|
if conflicts:
|
formatted = "\n".join(f"- {path}" for path in conflicts)
|
raise MediaTranscriptionError(
|
"输出文件已存在；默认不会覆盖。请改用新的输出目录或先人工处理已有结果：\n"
|
f"{formatted}"
|
)
|
|
|
def _probe_first_audio_stream(
|
video: Path,
|
ffprobe: str,
|
runner: ProcessRunner,
|
) -> None:
|
command = [
|
ffprobe,
|
"-v",
|
"error",
|
"-select_streams",
|
"a:0",
|
"-show_entries",
|
"stream=index",
|
"-of",
|
"csv=p=0",
|
str(video),
|
]
|
result = runner(command)
|
if result.returncode != 0:
|
raise MediaTranscriptionError(
|
f"FFprobe 无法检查视频音轨：{_error_detail(result)}"
|
)
|
if not result.stdout.strip():
|
raise MediaTranscriptionError(f"视频不包含可用音轨：{video}")
|
|
|
def _extract_first_audio_stream(
|
video: Path,
|
audio_path: Path,
|
ffmpeg: str,
|
runner: ProcessRunner,
|
) -> None:
|
command = [
|
ffmpeg,
|
"-hide_banner",
|
"-loglevel",
|
"error",
|
"-nostdin",
|
"-n",
|
"-i",
|
str(video),
|
"-map",
|
"0:a:0",
|
"-vn",
|
"-ac",
|
"1",
|
"-ar",
|
"16000",
|
"-c:a",
|
"flac",
|
str(audio_path),
|
]
|
result = runner(command)
|
if result.returncode != 0:
|
raise MediaTranscriptionError(
|
f"FFmpeg 提取第一条音轨失败：{_error_detail(result)}"
|
)
|
if not audio_path.is_file() or audio_path.stat().st_size == 0:
|
raise MediaTranscriptionError("FFmpeg 未生成有效的 FLAC 音频文件。")
|
|
|
def _probe_audio_duration(
|
audio_path: Path,
|
ffprobe: str,
|
runner: ProcessRunner,
|
) -> float:
|
command = [
|
ffprobe,
|
"-v",
|
"error",
|
"-show_entries",
|
"format=duration",
|
"-of",
|
"default=noprint_wrappers=1:nokey=1",
|
str(audio_path),
|
]
|
result = runner(command)
|
if result.returncode != 0:
|
raise MediaTranscriptionError(
|
f"FFprobe 无法读取 FLAC 时长：{_error_detail(result)}"
|
)
|
try:
|
duration = float(result.stdout.strip())
|
except ValueError as exc:
|
raise MediaTranscriptionError(
|
f"FFprobe 返回了无效的 FLAC 时长：{result.stdout.strip() or '空'}"
|
) from exc
|
if not math.isfinite(duration) or duration <= 0:
|
raise MediaTranscriptionError(f"FLAC 时长必须为有限正数，实际为：{duration}")
|
return duration
|
|
|
def _plan_audio_chunks(duration: float) -> tuple[AudioChunk, ...]:
|
if not math.isfinite(duration) or duration <= SHORT_AUDIO_LIMIT_SECONDS:
|
return ()
|
|
total = int(math.ceil(duration / CHUNK_CORE_SECONDS))
|
chunks: list[AudioChunk] = []
|
for offset in range(total):
|
core_start = offset * CHUNK_CORE_SECONDS
|
core_end = min(duration, core_start + CHUNK_CORE_SECONDS)
|
chunks.append(
|
AudioChunk(
|
index=offset + 1,
|
total=total,
|
core_start=core_start,
|
core_end=core_end,
|
audio_start=max(0.0, core_start - CHUNK_OVERLAP_SECONDS),
|
audio_end=min(duration, core_end + CHUNK_OVERLAP_SECONDS),
|
)
|
)
|
return tuple(chunks)
|
|
|
def _extract_audio_chunk(
|
full_audio: Path,
|
chunk_path: Path,
|
chunk: AudioChunk,
|
ffmpeg: str,
|
runner: ProcessRunner,
|
) -> None:
|
if chunk_path.exists():
|
raise MediaTranscriptionError(f"临时音频块已存在，拒绝覆盖：{chunk_path}")
|
command = [
|
ffmpeg,
|
"-hide_banner",
|
"-loglevel",
|
"error",
|
"-nostdin",
|
"-n",
|
"-ss",
|
f"{chunk.audio_start:.6f}",
|
"-i",
|
str(full_audio),
|
"-t",
|
f"{chunk.audio_duration:.6f}",
|
"-ac",
|
"1",
|
"-ar",
|
"16000",
|
"-c:a",
|
"flac",
|
str(chunk_path),
|
]
|
result = runner(command)
|
if result.returncode != 0:
|
raise MediaTranscriptionError(
|
f"FFmpeg 提取转写分块 {chunk.index}/{chunk.total} 失败："
|
f"{_error_detail(result)}"
|
)
|
if not chunk_path.is_file() or chunk_path.stat().st_size == 0:
|
raise MediaTranscriptionError(
|
f"FFmpeg 未生成有效的转写分块 {chunk.index}/{chunk.total}。"
|
)
|
|
|
def _load_cuda_model() -> tuple[Any, RuntimeInfo]:
|
try:
|
import ctranslate2
|
from faster_whisper import WhisperModel
|
except ImportError as exc:
|
raise MediaTranscriptionError(
|
"缺少 faster-whisper/CTranslate2，请在项目隔离 Conda 环境中运行。"
|
) from exc
|
|
try:
|
device_count = int(ctranslate2.get_cuda_device_count())
|
compute_types = tuple(
|
sorted(str(item) for item in ctranslate2.get_supported_compute_types(DEVICE))
|
)
|
except Exception as exc:
|
raise MediaTranscriptionError(
|
f"CUDA GPU 自检失败；不会降级到 CPU。原始错误：{exc}"
|
) from exc
|
|
if device_count < 1:
|
raise MediaTranscriptionError("未检测到可用 CUDA GPU；不会降级到 CPU。")
|
if COMPUTE_TYPE not in compute_types:
|
raise MediaTranscriptionError(
|
"当前 CUDA GPU 不支持 float16；不会降级到 CPU。"
|
f"检测到的计算类型：{', '.join(compute_types) or '无'}"
|
)
|
|
try:
|
model = WhisperModel(
|
MODEL_NAME,
|
device=DEVICE,
|
compute_type=COMPUTE_TYPE,
|
)
|
except Exception as exc:
|
raise MediaTranscriptionError(
|
"large-v3 模型加载失败；请检查模型缓存/网络以及 CUDA 12、cuBLAS、"
|
f"cuDNN 9。不会降级到 CPU。原始错误：{exc}"
|
) from exc
|
|
return model, RuntimeInfo(device_count, compute_types)
|
|
|
def _collect_transcript(
|
model: Any,
|
audio_path: Path,
|
language: str | None,
|
) -> TranscriptResult:
|
options: dict[str, Any] = {"vad_filter": True}
|
if language:
|
options["language"] = language
|
|
try:
|
segment_iterator, info = model.transcribe(str(audio_path), **options)
|
segments = tuple(_normalize_segments(segment_iterator))
|
except Exception as exc:
|
raise MediaTranscriptionError(
|
"GPU 语音转写失败；请检查 CUDA/cuDNN/显存和输入音频。"
|
f"不会降级到 CPU。原始错误：{exc}"
|
) from exc
|
|
detected_language = getattr(info, "language", None)
|
probability = getattr(info, "language_probability", None)
|
if probability is not None:
|
probability = float(probability)
|
return TranscriptResult(detected_language, probability, segments)
|
|
|
def _normalize_segments(segments: Iterable[Any]) -> Iterable[TranscriptSegment]:
|
for segment in segments:
|
text = str(segment.text).strip()
|
if not text:
|
continue
|
start = max(0.0, float(segment.start))
|
end = max(start, float(segment.end))
|
yield TranscriptSegment(start=start, end=end, text=text)
|
|
|
def _project_chunk_segments(
|
segments: Iterable[TranscriptSegment],
|
chunk: AudioChunk,
|
audio_duration: float,
|
) -> tuple[TranscriptSegment, ...]:
|
projected: list[TranscriptSegment] = []
|
for segment in segments:
|
start = min(audio_duration, max(0.0, chunk.audio_start + segment.start))
|
end = min(audio_duration, max(start, chunk.audio_start + segment.end))
|
midpoint = (start + end) / 2.0
|
belongs = chunk.core_start <= midpoint < chunk.core_end
|
if chunk.index == chunk.total:
|
belongs = chunk.core_start <= midpoint <= chunk.core_end
|
if belongs:
|
projected.append(TranscriptSegment(start=start, end=end, text=segment.text))
|
return tuple(projected)
|
|
|
def _collect_chunked_transcript(
|
model: Any,
|
full_audio: Path,
|
audio_duration: float,
|
language: str | None,
|
ffmpeg: str,
|
runner: ProcessRunner,
|
) -> TranscriptResult:
|
chunks = _plan_audio_chunks(audio_duration)
|
if not chunks:
|
raise MediaTranscriptionError("内部错误：短音频不应进入分块转写路径。")
|
|
chunk_path = full_audio.parent / ".transcribe-chunk.flac"
|
active_language = language
|
result_language = language
|
language_probability: float | None = None
|
projected: list[TranscriptSegment] = []
|
|
for chunk in chunks:
|
print(f"正在转写分块 {chunk.index}/{chunk.total}", flush=True)
|
try:
|
_extract_audio_chunk(full_audio, chunk_path, chunk, ffmpeg, runner)
|
chunk_result = _collect_transcript(model, chunk_path, active_language)
|
if chunk.index == 1:
|
result_language = chunk_result.language or language
|
language_probability = chunk_result.language_probability
|
if language is None:
|
if not chunk_result.language:
|
raise MediaTranscriptionError(
|
"首个音频块未返回检测语言，无法为后续块固定语言。"
|
)
|
active_language = chunk_result.language
|
projected.extend(
|
_project_chunk_segments(chunk_result.segments, chunk, audio_duration)
|
)
|
finally:
|
chunk_path.unlink(missing_ok=True)
|
|
projected.sort(key=lambda item: (item.start, item.end, item.text))
|
return TranscriptResult(
|
language=result_language,
|
language_probability=language_probability,
|
segments=tuple(projected),
|
)
|
|
|
def _format_txt_timestamp(seconds: float) -> str:
|
total_seconds = max(0, int(seconds))
|
hours, remainder = divmod(total_seconds, 3600)
|
minutes, secs = divmod(remainder, 60)
|
return f"{hours:02d}:{minutes:02d}:{secs:02d}"
|
|
|
def _format_srt_timestamp(seconds: float) -> str:
|
total_milliseconds = max(0, int(round(seconds * 1000)))
|
hours, remainder = divmod(total_milliseconds, 3_600_000)
|
minutes, remainder = divmod(remainder, 60_000)
|
secs, milliseconds = divmod(remainder, 1000)
|
return f"{hours:02d}:{minutes:02d}:{secs:02d},{milliseconds:03d}"
|
|
|
def _write_transcript_files(
|
video: Path,
|
staged: OutputPaths,
|
result: TranscriptResult,
|
) -> None:
|
text_content = "\n\n".join(
|
f"[{_format_txt_timestamp(segment.start)}] {segment.text}"
|
for segment in result.segments
|
)
|
if text_content:
|
text_content += "\n"
|
staged.text.write_text(text_content, encoding="utf-8")
|
|
srt_blocks = []
|
for index, segment in enumerate(result.segments, start=1):
|
srt_blocks.append(
|
"\n".join(
|
[
|
str(index),
|
f"{_format_srt_timestamp(segment.start)} --> "
|
f"{_format_srt_timestamp(segment.end)}",
|
segment.text,
|
]
|
)
|
)
|
srt_content = "\n\n".join(srt_blocks)
|
if srt_content:
|
srt_content += "\n"
|
staged.srt.write_text(srt_content, encoding="utf-8")
|
|
payload = {
|
"source": video.name,
|
"audio": staged.audio.name,
|
"language": result.language,
|
"language_probability": result.language_probability,
|
"model": MODEL_NAME,
|
"device": DEVICE,
|
"compute_type": COMPUTE_TYPE,
|
"vad_filter": True,
|
"segments": [
|
{
|
"start": segment.start,
|
"end": segment.end,
|
"text": segment.text,
|
}
|
for segment in result.segments
|
],
|
}
|
staged.json.write_text(
|
json.dumps(payload, ensure_ascii=False, indent=2) + "\n",
|
encoding="utf-8",
|
)
|
|
|
def _staged_paths(staging_dir: Path, final: OutputPaths) -> OutputPaths:
|
return OutputPaths(
|
directory=staging_dir,
|
audio=staging_dir / final.audio.name,
|
text=staging_dir / final.text.name,
|
srt=staging_dir / final.srt.name,
|
json=staging_dir / final.json.name,
|
)
|
|
|
def _commit_outputs(staged: OutputPaths, final: OutputPaths) -> None:
|
_ensure_output_available(final)
|
created_directory = False
|
moved: list[Path] = []
|
try:
|
if not final.directory.exists():
|
final.directory.mkdir()
|
created_directory = True
|
for source, destination in zip(staged.files, final.files, strict=True):
|
try:
|
source.replace(destination)
|
except BaseException:
|
if destination.exists() and not source.exists():
|
moved.append(destination)
|
raise
|
else:
|
moved.append(destination)
|
except BaseException as exc:
|
rollback_failures: list[str] = []
|
for path in reversed(moved):
|
try:
|
path.unlink(missing_ok=True)
|
except BaseException as cleanup_exc:
|
rollback_failures.append(f"{path}: {cleanup_exc}")
|
if created_directory:
|
try:
|
final.directory.rmdir()
|
except BaseException as cleanup_exc:
|
rollback_failures.append(f"{final.directory}: {cleanup_exc}")
|
|
if not isinstance(exc, Exception):
|
if rollback_failures:
|
try:
|
setattr(
|
exc,
|
"_media_commit_rollback_failures",
|
tuple(rollback_failures),
|
)
|
except Exception:
|
pass
|
raise
|
|
rollback_detail = ""
|
if rollback_failures:
|
rollback_detail = "；回滚不完整：" + "；".join(rollback_failures)
|
raise MediaTranscriptionError(
|
f"提交输出文件失败，已回滚已写结果：{exc}{rollback_detail}"
|
) from exc
|
|
|
def transcribe_video(
|
video_path: str | Path,
|
output_dir: str | Path | None = None,
|
language: str | None = None,
|
*,
|
runner: ProcessRunner | None = None,
|
model_loader: ModelLoader | None = None,
|
) -> tuple[OutputPaths, RuntimeInfo]:
|
runner = runner or _run_process
|
model_loader = model_loader or _load_cuda_model
|
language = language.strip().lower() if language else None
|
|
video = _resolve_video(video_path)
|
outputs = _build_output_paths(video, output_dir)
|
_ensure_output_available(outputs)
|
ffprobe = _require_executable("ffprobe")
|
ffmpeg = _require_executable("ffmpeg")
|
_probe_first_audio_stream(video, ffprobe, runner)
|
|
outputs.directory.parent.mkdir(parents=True, exist_ok=True)
|
staging_dir = Path(
|
tempfile.mkdtemp(
|
prefix=f".{video.stem}.transcribe-",
|
dir=str(outputs.directory.parent),
|
)
|
)
|
staged = _staged_paths(staging_dir, outputs)
|
|
try:
|
print("正在提取第一条音轨（16 kHz、单声道 FLAC）……")
|
_extract_first_audio_stream(video, staged.audio, ffmpeg, runner)
|
audio_duration = _probe_audio_duration(staged.audio, ffprobe, runner)
|
print("正在加载 large-v3（CUDA、float16）……")
|
model, runtime = model_loader()
|
print("正在使用 GPU 和 VAD 转写……")
|
if audio_duration <= SHORT_AUDIO_LIMIT_SECONDS:
|
transcript = _collect_transcript(model, staged.audio, language)
|
else:
|
transcript = _collect_chunked_transcript(
|
model,
|
staged.audio,
|
audio_duration,
|
language,
|
ffmpeg,
|
runner,
|
)
|
_write_transcript_files(video, staged, transcript)
|
_commit_outputs(staged, outputs)
|
finally:
|
shutil.rmtree(staging_dir, ignore_errors=True)
|
|
return outputs, runtime
|
|
|
def _build_parser() -> argparse.ArgumentParser:
|
parser = argparse.ArgumentParser(
|
description="用 FFmpeg 和 faster-whisper GPU 转写单个本地视频。"
|
)
|
parser.add_argument("video", help="本地视频文件路径")
|
parser.add_argument("--output", help="可选输出目录；默认在视频旁创建 <文件名>.transcript")
|
parser.add_argument("--language", help="可选语言代码，例如 zh；默认自动检测")
|
return parser
|
|
|
def main(argv: Sequence[str] | None = None) -> int:
|
args = _build_parser().parse_args(argv)
|
try:
|
outputs, runtime = transcribe_video(args.video, args.output, args.language)
|
except MediaTranscriptionError as exc:
|
print(f"错误：{exc}", file=sys.stderr)
|
return 1
|
except KeyboardInterrupt as exc:
|
rollback_failures = getattr(exc, "_media_commit_rollback_failures", ())
|
if rollback_failures:
|
detail = "\n".join(f"- {item}" for item in rollback_failures)
|
print(
|
"错误：用户中断；正式输出回滚不完整，请人工检查以下路径：\n"
|
f"{detail}",
|
file=sys.stderr,
|
)
|
else:
|
print("错误：用户中断，未提交输出文件。", file=sys.stderr)
|
return 130
|
|
print(f"完成：{outputs.directory}")
|
print(
|
"GPU 自检："
|
f"CUDA 设备数={runtime.cuda_device_count}，"
|
f"计算类型={','.join(runtime.cuda_compute_types)}"
|
)
|
return 0
|
|
|
if __name__ == "__main__":
|
raise SystemExit(main())
|
