whisper-transcribe/scripts/whisper_transcribe.py

612 lines
24 KiB
Python
Raw Permalink Blame History

This file contains ambiguous Unicode characters

This file contains Unicode characters that might be confused with other characters. If you think that this is intentional, you can safely ignore this warning. Use the Escape button to reveal them.

#!/usr/bin/env python3
"""
MLX Whisper enhanced transcription with precise subtitle timing.
Uses mlx-whisper (Apple Silicon optimized) + word-level timestamps + smart split.
Usage:
python3 whisper_transcribe.py <audio_file> --initial-prompt "<提示词>" [选项]
Options:
--model MODEL mlx-whisper model: turbo/small/medium (default: turbo)
--language LANG Language code (ja/en/zh/...)
--output-dir DIR Output directory (default: same as input)
--output-format FMT Output format: txt/srt/json/all (default: txt)
--no-clean Disable hallucination cleanup
--no-word-ts Disable word-level timestamps (revert to segment-level)
--no-smart-split Disable smart split on pauses (default: on)
--no-vad-filter Disable silence/VAD filtering (default: on)
--max-line-duration Max duration per subtitle line in seconds (default: 5.0)
--initial-prompt TEXT Initial prompt for context guidance (required)
--verbose Show progress
"""
import argparse
import json
import os
import re
import subprocess
import sys
import time
from contextlib import contextmanager
from pathlib import Path
from typing import Optional
import tempfile as tf_module
import shutil
# mlx-whisper model repo mapping
MLX_MODELS = {
"turbo": "mlx-community/whisper-large-v3-turbo",
"large-v3-turbo": "mlx-community/whisper-large-v3-turbo",
"small": "mlx-community/whisper-small-mlx",
"medium": "mlx-community/whisper-medium-mlx",
}
# ── Audio Preprocessing ──────────────────────────────────────────────────────
def preprocess_audio(audio_path: str, output_dir: str) -> str:
"""Preprocess audio: downsample to 16kHz mono + loudness normalization."""
out_path = os.path.join(output_dir, "_preprocessed.wav")
cmd = [
"ffmpeg", "-y", "-i", str(audio_path),
"-af", "loudnorm=I=-16:TP=-1.5:LRA=11",
"-ar", "16000", "-ac", "1",
"-f", "wav", out_path,
]
result = subprocess.run(cmd, capture_output=True, text=True)
if result.returncode != 0:
raise RuntimeError(f"ffmpeg preprocessing failed: {result.stderr}")
return out_path
def get_audio_duration(audio_path: str) -> float:
"""Get audio duration in seconds using ffprobe."""
cmd = [
"ffprobe", "-v", "quiet", "-print_format", "json",
"-show_format", str(audio_path)
]
result = subprocess.run(cmd, capture_output=True, text=True)
if result.returncode != 0:
raise RuntimeError(f"ffprobe failed: {result.stderr}")
info = json.loads(result.stdout)
return float(info["format"]["duration"])
def detect_silence(audio_path: str, noise_db: int = -40, min_duration: float = 0.3,
audio_duration: float = None) -> list[tuple[float, float]]:
"""Detect silent regions using ffmpeg silencedetect.
Returns list of (start, end) tuples for silent segments.
If audio ends in silence, ffmpeg emits silence_start without silence_end;
we use audio_duration to close the last unclosed silence region."""
cmd = [
"ffmpeg", "-i", str(audio_path),
"-af", f"silencedetect=noise={noise_db}dB:d={min_duration}",
"-f", "null", "-"
]
result = subprocess.run(cmd, capture_output=True, text=True)
starts = re.findall(r"silence_start: ([\d.]+)", result.stderr)
ends = re.findall(r"silence_end: ([\d.]+)", result.stderr)
starts = [float(s) for s in starts]
ends = [float(e) for e in ends]
silences = []
# Pair starts with ends; if there's an unmatched start at the end,
# close it with audio_duration
for i, s in enumerate(starts):
if i < len(ends):
silences.append((s, ends[i]))
elif audio_duration:
# Unclosed silence at end of audio
silences.append((s, audio_duration))
return silences
# ── Transcription ────────────────────────────────────────────────────────────
def transcribe_file(audio_path: str, model_repo: str, language: str,
initial_prompt: Optional[str] = None,
word_timestamps: bool = True,
verbose: bool = False) -> dict:
"""Transcribe a single audio file with mlx-whisper."""
from mlx_whisper import transcribe
options = {
"path_or_hf_repo": model_repo,
"language": language,
"verbose": verbose,
"word_timestamps": word_timestamps,
}
if initial_prompt:
options["initial_prompt"] = initial_prompt
return transcribe(audio_path, **options)
# ── Post-processing: Smart Split & Glue Fix ──────────────────────────────────
def smart_split_segments(segments: list[dict], max_duration: float = 5.0,
min_pause: float = 0.3) -> list[dict]:
"""Split long segments on natural pauses using word-level timestamps.
Fixes subtitle 'glue' where one line sticks on screen too long."""
result = []
for seg in segments:
words = seg.get("words", [])
if not words:
result.append(seg)
continue
# If segment is short enough, keep as-is
if seg["end"] - seg["start"] <= max_duration:
result.append(seg)
continue
# Split on pauses between words
current_words = []
current_start = words[0]["start"]
for i, w in enumerate(words):
current_words.append(w)
# Check for pause after this word
pause = 0.0
if i < len(words) - 1:
pause = words[i + 1]["start"] - w["end"]
# Split condition: pause long enough OR current chunk exceeds max_duration
chunk_duration = w["end"] - current_start
should_split = (pause >= min_pause and chunk_duration >= 1.5) or \
chunk_duration >= max_duration
if should_split and current_words:
text = "".join(wd["word"] for wd in current_words).strip()
if text:
result.append({
"start": current_start,
"end": w["end"],
"text": text,
"words": current_words[:],
})
current_words = []
if i < len(words) - 1:
current_start = words[i + 1]["start"]
# Flush remaining words
if current_words:
text = "".join(wd["word"] for wd in current_words).strip()
if text:
result.append({
"start": current_start,
"end": current_words[-1]["end"],
"text": text,
"words": current_words[:],
})
return result
def remove_vad_gaps(segments: list[dict], silences: list[tuple[float, float]],
min_gap: float = 0.5) -> list[dict]:
"""Remove segments that fall entirely within detected silent regions.
Prevents 'ghost subtitles' during music-only sections."""
if not silences:
return segments
def in_silence(start: float, end: float) -> bool:
for s_start, s_end in silences:
# If segment is mostly inside a silence region
overlap_start = max(start, s_start)
overlap_end = min(end, s_end)
if overlap_end > overlap_start:
overlap = overlap_end - overlap_start
seg_len = end - start
if seg_len <= 0:
continue # Skip zero-length segments
if overlap / seg_len > 0.7: # 70% or more in silence
return True
return False
filtered = []
for seg in segments:
if not in_silence(seg["start"], seg["end"]):
filtered.append(seg)
return filtered
def clean_hallucinations(segments: list[dict]) -> list[dict]:
"""Remove trailing hallucinated repetitions."""
if not segments:
return segments
texts = [s["text"] for s in segments]
cut_point = len(texts)
seen_repeats = 0
max_repeats = 5
for i in range(len(texts) - 1, 0, -1):
if texts[i] == texts[i - 1] and texts[i].strip():
seen_repeats += 1
if seen_repeats >= max_repeats:
cut_point = i
break
else:
seen_repeats = 0
filler_pattern = re.compile(r'^(あ+|啊+|um+|uh+|er+|ah+)$', re.IGNORECASE)
filler_count = 0
for i in range(len(texts) - 1, max(cut_point - 10, 0), -1):
if filler_pattern.match(texts[i].strip()):
filler_count += 1
else:
break
if filler_count > 3:
cut_point = len(texts) - filler_count
return segments[:cut_point]
def is_phrase_repetition(text: str) -> bool:
"""检测短语/词组级别的幻听死循环。
只对大段重复敏感:精确周期需 ≥10 次重复n-gram 需 ≥30 字符。
短文本如「うんうんうん」「うふふふ」正常放行。"""
t = text.strip()
if len(t) < 6:
return False
# 精确周期性检测: 如果 t 由某个子串重复多次构成
pos = (t + t).find(t, 1)
if pos != -1 and pos < len(t):
# 至少重复 10 次以上才算循环(避免 "うん"×3 这种正常相槌)
if len(t) // pos >= 10:
return True
# 宽松检测:前后两半完全相同或包含关系(仅对较长文本)
if len(t) >= 16:
mid = len(t) // 2
if t[:mid] == t[mid:] or (mid >= 3 and t[:mid] in t[mid:]):
return True
# 近似周期检测:对付长度不对齐的循环(如 "ABC"*55 + "AB"
# n-gram 去重率极低 → 文本由少数短语反复拼接
if len(t) >= 30:
for n in (3, 4, 5):
grams = [t[i:i + n] for i in range(len(t) - n + 1)]
unique_ratio = len(set(grams)) / len(grams)
if unique_ratio < 0.15:
return True
return False
def fix_repetition_loops(segments: list[dict], audio_path: str,
model_repo: str, language: str,
audio_duration: float,
initial_prompt: Optional[str] = None,
char_threshold: float = 0.7,
silences: Optional[list[tuple[float, float]]] = None) -> list[dict]:
"""检测并重新转写 Whisper 陷入死循环或长空段的异常片段。
Detection modes:
1. Single-char loop — one character dominates (e.g. '秀秀秀秀秀...')
2. Phrase loop — text built from a repeating substring
3. Empty segment — text empty AND segment ≥15s AND not in VAD silence
If retry produces no valid text, the segment range is deleted.
Traverses ranges in REVERSE to avoid index shifting from list mutations.
"""
# Helper: check if a time range overlaps any VAD silence region
def _in_silence(seg_start: float, seg_end: float) -> bool:
if not silences:
return False
for s_start, s_end in silences:
if seg_start < s_end and seg_end > s_start:
return True
return False
# 1. Detect problem segments
bad_indices = []
for i, seg in enumerate(segments):
text = seg.get("text", "").strip()
# Mode A1: single-char loop (only for very long repetitions)
if len(text) >= 20:
most_common = max(set(text), key=text.count)
if text.count(most_common) / len(text) > char_threshold:
bad_indices.append(i)
continue
# Mode A2: phrase-level loop
if is_phrase_repetition(text):
bad_indices.append(i)
continue
# Mode B: long empty segment not explainable by VAD silence
seg_len = seg["end"] - seg["start"]
if not text and seg_len >= 15.0 and not _in_silence(seg["start"], seg["end"]):
bad_indices.append(i)
if not bad_indices:
return segments
# 2. Group consecutive indices into continuous ranges
ranges = []
start_i = bad_indices[0]
end_i = bad_indices[0]
for i in bad_indices[1:]:
if i == end_i + 1:
end_i = i
else:
ranges.append((start_i, end_i))
start_i = end_i = i
ranges.append((start_i, end_i))
# 3. Re-transcribe each range with a cold start
# CRITICAL: reversed() prevents index shifting when list length changes
from mlx_whisper import transcribe
result = list(segments)
total_replaced = 0
for seg_start_idx, seg_end_idx in reversed(ranges):
t_start = max(0, result[seg_start_idx]["start"] - 0.3)
t_end = min(result[seg_end_idx]["end"] + 0.3, audio_duration)
print(f" Problem segment at {t_start:.1f}s-{t_end:.1f}s, "
f"re-transcribing ({seg_end_idx - seg_start_idx + 1} segs)...")
tmp_dir = tf_module.mkdtemp()
try:
seg_path = os.path.join(tmp_dir, "_retry.wav")
subprocess.run([
"ffmpeg", "-y", "-i", audio_path,
"-ss", str(t_start),
"-to", str(t_end),
"-af", "loudnorm=I=-16:TP=-1.5:LRA=11",
"-ar", "16000", "-ac", "1",
"-f", "wav", seg_path,
], capture_output=True, text=True, check=True)
# Cold-start transcribe (no preceding context)
retry_opts = {
"path_or_hf_repo": model_repo,
"language": language,
"verbose": False,
"word_timestamps": True,
}
if initial_prompt:
retry_opts["initial_prompt"] = initial_prompt
retry_result = transcribe(seg_path, **retry_opts)
# Build replacement segments with absolute timestamps
new_segs = []
for s in retry_result.get("segments", []):
text_clean = s["text"].strip()
# Filter: no empty text, no repeated loops in retry
if not text_clean:
continue
if len(text_clean) >= 5:
mc = max(set(text_clean), key=text_clean.count)
if text_clean.count(mc) / len(text_clean) > char_threshold:
continue
if is_phrase_repetition(text_clean):
continue
# Density check: impossible char rate → hallucination
# Normal Japanese speech is ~7 chars/s; 20 is very generous
seg_dur = s["end"] - s["start"]
if seg_dur > 0 and len(text_clean) > seg_dur * 20:
continue
new_segs.append({
"start": s["start"] + t_start,
"end": s["end"] + t_start,
"text": text_clean,
"words": [
{"start": w["start"] + t_start,
"end": w["end"] + t_start,
"word": w["word"]}
for w in s.get("words", [])
]
})
if new_segs:
result[seg_start_idx:seg_end_idx + 1] = new_segs
total_replaced += (seg_end_idx - seg_start_idx + 1)
print(f" Replaced with {len(new_segs)} corrected segments")
else:
# Retry produced nothing useful — delete (silence/noise/hallucination)
del result[seg_start_idx:seg_end_idx + 1]
print(f" Retry empty, removed {seg_end_idx - seg_start_idx + 1} segment(s)")
except Exception as e:
print(f" Warning: retry failed ({e}), keeping original segments")
finally:
shutil.rmtree(tmp_dir, ignore_errors=True)
if total_replaced:
print(f" Repetition fix: replaced {total_replaced} hallucinated segments")
return result
# ── Output Formatting ────────────────────────────────────────────────────────
def format_timestamp_srt(seconds: float) -> str:
"""Format seconds to SRT timestamp: HH:MM:SS,mmm
Uses round() to avoid floating-point precision errors where
1.001 % 1 = 0.000999... → int() truncates to 000 instead of 001.
Clamps negative values to 0 (Whisper shouldn't return them, but be safe)."""
if seconds < 0:
seconds = 0.0
total_ms = int(round(seconds * 1000))
h = total_ms // 3600000
m = (total_ms % 3600000) // 60000
s = (total_ms % 60000) // 1000
ms = total_ms % 1000
return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"
def write_output(segments: list[dict], output_path: str, fmt: str):
"""Write transcription output in specified format."""
base = Path(output_path).with_suffix("")
if fmt in ("txt", "all"):
txt_path = str(base) + ".txt"
with open(txt_path, "w", encoding="utf-8") as f:
for seg in segments:
f.write(seg["text"] + "\n")
print(f" Written: {txt_path}")
if fmt in ("srt", "all"):
srt_path = str(base) + ".srt"
with open(srt_path, "w", encoding="utf-8") as f:
for i, seg in enumerate(segments, 1):
f.write(f"{i}\n")
f.write(f"{format_timestamp_srt(seg['start'])} --> "
f"{format_timestamp_srt(seg['end'])}\n")
f.write(f"{seg['text']}\n\n")
print(f" Written: {srt_path}")
if fmt in ("json", "all"):
json_path = str(base) + ".json"
with open(json_path, "w", encoding="utf-8") as f:
json.dump(segments, f, ensure_ascii=False, indent=2)
print(f" Written: {json_path}")
# ── Main ─────────────────────────────────────────────────────────────────────
@contextmanager
def tempfile_directory():
d = tf_module.mkdtemp()
try:
yield d
finally:
shutil.rmtree(d, ignore_errors=True)
def main():
parser = argparse.ArgumentParser(
description="MLX Whisper enhanced transcription with precise subtitles")
parser.add_argument("audio", help="Input audio file path")
parser.add_argument("--model", default="turbo",
choices=["turbo", "large-v3-turbo", "small", "medium"],
help="mlx-whisper model (default: turbo)")
parser.add_argument("--language", default="ja",
help="Language code (default: ja)")
parser.add_argument("--output-dir", default=None,
help="Output directory (default: same as input)")
parser.add_argument("--output-format", default="txt",
choices=["txt", "srt", "json", "all"],
help="Output format (default: txt)")
parser.add_argument("--no-clean", action="store_true",
help="Disable hallucination cleanup")
parser.add_argument("--no-word-ts", action="store_true",
help="Disable word-level timestamps")
parser.add_argument("--no-smart-split", action="store_true",
help="Disable smart split on pauses")
parser.add_argument("--no-vad-filter", action="store_true",
help="Disable silence/VAD filtering")
parser.add_argument("--max-line-duration", type=float, default=5.0,
help="Max seconds per subtitle line (default: 5.0)")
parser.add_argument("--initial-prompt", required=True,
help="Initial prompt for context guidance (required)")
parser.add_argument("--verbose", action="store_true",
help="Show progress")
args = parser.parse_args()
audio_path = Path(args.audio).resolve()
if not audio_path.exists():
print(f"Error: File not found: {audio_path}", file=sys.stderr)
sys.exit(1)
output_dir = Path(args.output_dir) if args.output_dir else audio_path.parent
output_dir.mkdir(parents=True, exist_ok=True)
output_name = audio_path.stem
model_repo = MLX_MODELS[args.model]
# Get duration
duration = get_audio_duration(str(audio_path))
duration_min = duration / 60
print(f"Audio: {audio_path.name} ({duration_min:.1f} min)")
print(f"Model: {args.model} ({model_repo})")
# ── Preprocessing ──────────────────────────────────────────────────────
print("Preprocessing: 16kHz mono + loudness normalization...")
with tempfile_directory() as preprocess_dir:
preprocessed = preprocess_audio(str(audio_path), preprocess_dir)
print(f" Preprocessed: {Path(preprocessed).name}")
# Optional: detect silence regions for VAD filtering
silences = []
if not args.no_vad_filter:
print("Detecting silent regions...")
silences = detect_silence(preprocessed, noise_db=-40, min_duration=0.5,
audio_duration=duration)
print(f" Found {len(silences)} silent regions")
# ── Transcribe (mlx-whisper handles long audio internally) ──────────
print(f"Transcribing with mlx-whisper ({args.model})...")
t0 = time.time()
result = transcribe_file(
preprocessed, model_repo, args.language,
args.initial_prompt,
word_timestamps=not args.no_word_ts,
verbose=args.verbose)
elapsed = time.time() - t0
segments_data = result.get("segments", [])
merged = [
{"start": s["start"], "end": s["end"],
"text": s["text"].strip(),
"words": [
{"start": w["start"], "end": w["end"], "word": w["word"]}
for w in s.get("words", [])
]}
for s in segments_data
]
print(f" Transcribed {len(merged)} segments in {elapsed:.1f}s")
# ── Post-processing ────────────────────────────────────────────────────
# 1. VAD filter: remove ghost subtitles in silent regions
if not args.no_vad_filter and silences:
before = len(merged)
merged = remove_vad_gaps(merged, silences)
removed = before - len(merged)
if removed > 0:
print(f" VAD filter: removed {removed} segments in silent regions")
# 2. Smart split: break long glued segments on natural pauses
if not args.no_smart_split and not args.no_word_ts:
before = len(merged)
merged = smart_split_segments(merged, max_duration=args.max_line_duration)
if len(merged) > before:
print(f" Smart split: {before}{len(merged)} segments")
# 3. Hallucination cleanup
if not args.no_clean:
before = len(merged)
merged = clean_hallucinations(merged)
removed = before - len(merged)
if removed > 0:
print(f" Cleaned: removed {removed} hallucinated segments")
# 4. Repetition loop detection & local re-transcription
if not args.no_clean:
merged = fix_repetition_loops(
merged, str(audio_path), model_repo,
args.language, duration, args.initial_prompt,
silences=silences)
# ── Output ─────────────────────────────────────────────────────────────
out_path = str(output_dir / output_name)
print("Output:")
write_output(merged, out_path, args.output_format)
print("Done!")
if __name__ == "__main__":
main()