469 lines
19 KiB
Python
Executable File
469 lines
19 KiB
Python
Executable File
#!/usr/bin/env -S uv run --script
|
|
# /// script
|
|
# dependencies = ["httpx"]
|
|
# ///
|
|
import argparse
|
|
import base64
|
|
import functools
|
|
import logging
|
|
import os
|
|
import re
|
|
import subprocess
|
|
import tempfile
|
|
import textwrap
|
|
import time
|
|
from collections.abc import Callable
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import httpx
|
|
|
|
RETRY_STATUS_CODES = {429, 500, 502, 503, 504, 524, 529}
|
|
MAX_ATTEMPTS = 5
|
|
MIN_SPLIT_DURATION = 1.0
|
|
MIN_CHUNK_DURATION = 0.1
|
|
DIALOGUE_LINE_LENGTH = 42
|
|
MAX_DIALOGUE_GAP = 1.0
|
|
MAX_DIALOGUE_DURATION = 7.0
|
|
LINE_LENGTH = 42
|
|
MAX_LINES = 2
|
|
# lowercase tokens (en, tr) whose trailing period does not end a sentence
|
|
ABBREVIATIONS = {"mr", "mrs", "ms", "dr", "prof", "e.g", "i.e", "vs", "etc", "doç", "av", "sn", "vb", "bkz"}
|
|
# Finds candidate sentence boundaries: end punctuation, optional closers, then whitespace.
|
|
# Matches in "He left. Then" -> ". " and in 'He said "no." Then' -> '." '.
|
|
SENTENCE_END = re.compile(
|
|
r"""
|
|
[.!?…]+ # one or more end marks: "." "?!" "..." or the ellipsis character
|
|
[\"'”’)\]]* # closing quotes or brackets that belong to the finished sentence
|
|
\s+ # whitespace after it; one or more, so double spaces work too
|
|
""",
|
|
re.VERBOSE,
|
|
)
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
@dataclass
|
|
class Segment:
|
|
"""A timed piece of transcript; times are in seconds, speaker is a diarization index if known."""
|
|
|
|
start: float
|
|
end: float
|
|
text: str
|
|
speaker: int | None = None
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
"""Parse command-line arguments."""
|
|
parser = argparse.ArgumentParser(
|
|
description="Transcribe a video into an SRT file using an OpenRouter speech-to-text model.",
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
|
)
|
|
parser.add_argument("video", type=Path, help="Input video (or audio) file")
|
|
parser.add_argument("--output", "-o", type=Path, help="Output SRT path (default: <video>.srt)")
|
|
parser.add_argument("--model", "-m", default="elevenlabs/scribe-v2", help="OpenRouter STT model id (must support diarization unless --no-diarize)")
|
|
parser.add_argument(
|
|
"--api-key",
|
|
default=os.getenv("OPENROUTER_API_KEY"),
|
|
help="OpenRouter API key (default: $OPENROUTER_API_KEY%(default).0s)", # %(default).0s keeps the key out of --help
|
|
)
|
|
parser.add_argument("--language", "-l", help="ISO-639-1 source language, e.g. en or de (auto-detect if omitted)")
|
|
parser.add_argument("--chunk-minutes", type=float, default=60.0, help="Audio chunk length in minutes (of sped-up audio)")
|
|
parser.add_argument(
|
|
"--speed-up",
|
|
type=float,
|
|
default=1.0,
|
|
help="Speed the audio up (1.0-2.0) before transcribing to reduce cost; timestamps are scaled back and multi-sentence cues are split",
|
|
)
|
|
parser.add_argument(
|
|
"--diarize",
|
|
action=argparse.BooleanOptionalAction,
|
|
default=True,
|
|
help="Label speakers (needs a supporting model, e.g. elevenlabs/scribe-v2); short speaker changes become '- ' dialogue cues",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def run_ffmpeg_extract(video: Path, directory: Path, chunk_seconds: int, speed_up: float) -> list[Path]:
|
|
"""Extract mono 16 kHz mp3 audio into `directory` as chunks of `chunk_seconds` (measured after speed-up).
|
|
|
|
Returns the chunk paths in order. `speed_up` > 1 applies a single ffmpeg `atempo` filter.
|
|
"""
|
|
filter_arguments = ["-filter:a", f"atempo={speed_up}"] if speed_up != 1.0 else []
|
|
subprocess.run(
|
|
[
|
|
"ffmpeg",
|
|
"-hide_banner",
|
|
"-loglevel",
|
|
"error",
|
|
"-i",
|
|
str(video),
|
|
"-vn",
|
|
"-ac",
|
|
"1",
|
|
"-ar",
|
|
"16000",
|
|
*filter_arguments,
|
|
"-c:a",
|
|
"libmp3lame",
|
|
"-b:a",
|
|
"48k",
|
|
"-f",
|
|
"segment",
|
|
"-segment_time",
|
|
str(chunk_seconds),
|
|
"-reset_timestamps",
|
|
"1",
|
|
str(directory / "chunk_%04d.mp3"),
|
|
],
|
|
check=True,
|
|
)
|
|
chunks = sorted(directory.glob("chunk_*.mp3"))
|
|
if not chunks:
|
|
raise RuntimeError(f"ffmpeg produced no audio from {video}")
|
|
return chunks
|
|
|
|
|
|
def get_duration(path: Path) -> float:
|
|
"""Return the media duration in seconds via ffprobe, or 0.0 if it is unreadable (e.g. an empty trailing chunk)."""
|
|
result = subprocess.run(
|
|
["ffprobe", "-v", "error", "-show_entries", "format=duration", "-of", "csv=p=0", str(path)],
|
|
check=True,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
try:
|
|
return max(float(result.stdout.strip()), 0.0)
|
|
except ValueError:
|
|
return 0.0
|
|
|
|
|
|
def retry(
|
|
attempts: int = MAX_ATTEMPTS,
|
|
exceptions: tuple[type[Exception], ...] = (httpx.TransportError, httpx.HTTPStatusError),
|
|
) -> Callable:
|
|
"""Decorator factory: call the function up to `attempts` times, sleeping 2**attempt seconds between tries.
|
|
|
|
Only `exceptions` trigger a retry; the last one is re-raised when attempts run out.
|
|
"""
|
|
|
|
def decorator(func: Callable) -> Callable:
|
|
@functools.wraps(func)
|
|
def wrapper(*args, **kwargs):
|
|
for attempt in range(1, attempts + 1):
|
|
try:
|
|
return func(*args, **kwargs)
|
|
except exceptions as error:
|
|
if attempt == attempts:
|
|
raise
|
|
logger.warning("%s failed (%s), retry %d/%d", func.__name__, error, attempt, attempts)
|
|
time.sleep(2**attempt)
|
|
|
|
return wrapper
|
|
|
|
return decorator
|
|
|
|
|
|
@retry()
|
|
def post_request(client: httpx.Client, payload: dict) -> httpx.Response:
|
|
"""POST to the transcription endpoint. Retryable statuses (429/5xx) raise so `@retry` repeats them;
|
|
other statuses, including 400, are returned for the caller to handle."""
|
|
response = client.post("/audio/transcriptions", json=payload)
|
|
if response.status_code in RETRY_STATUS_CODES:
|
|
response.raise_for_status()
|
|
return response
|
|
|
|
|
|
def transcribe_chunk(
|
|
client: httpx.Client,
|
|
path: Path,
|
|
duration: float,
|
|
model: str,
|
|
language: str | None,
|
|
diarize: bool,
|
|
) -> tuple[list[Segment], bool]:
|
|
"""Transcribe one audio chunk. Times in the result are relative to the chunk, in sped-up audio time.
|
|
|
|
Returns (segments, has_timestamps). If the model rejects `verbose_json` the call is repeated as plain
|
|
JSON and a single segment covering the whole chunk is returned with has_timestamps=False.
|
|
With `diarize`, speaker turns are built from word-level labels, falling back to segment-level labels;
|
|
a rejected diarize request exits instead of falling back.
|
|
"""
|
|
payload: dict = {
|
|
"model": model,
|
|
"input_audio": {"data": base64.b64encode(path.read_bytes()).decode(), "format": "mp3"},
|
|
"response_format": "verbose_json",
|
|
"timestamp_granularities": ["segment", "word"] if diarize else ["segment"],
|
|
}
|
|
if language:
|
|
payload["language"] = language
|
|
if diarize:
|
|
payload["diarize"] = True
|
|
response = post_request(client=client, payload=payload)
|
|
if response.status_code == 400 and diarize:
|
|
raise SystemExit(f"Model rejected diarization: {response.text[:300]}")
|
|
if response.status_code == 400:
|
|
logger.warning("Model rejected verbose_json (%s); falling back to a single cue per chunk", response.text[:200])
|
|
payload.pop("response_format")
|
|
payload.pop("timestamp_granularities")
|
|
response = post_request(client=client, payload=payload)
|
|
response.raise_for_status()
|
|
data = response.json()
|
|
if not data.get("segments") and not data.get("words") and not str(data.get("text", "")).strip():
|
|
raise RuntimeError(f"Empty transcription for {path.name} ({path.stat().st_size / 1e6:.1f} MB); response: {str(data)[:300]}")
|
|
if diarize:
|
|
turns = turns_from_words(words=data.get("words") or [])
|
|
if turns:
|
|
return turns, True
|
|
segments = [
|
|
Segment(
|
|
start=float(segment["start"]),
|
|
end=float(segment["end"]),
|
|
text=str(segment["text"]).strip(),
|
|
speaker=segment.get("speaker") if diarize else None,
|
|
)
|
|
for segment in data.get("segments") or []
|
|
if str(segment.get("text", "")).strip()
|
|
]
|
|
if diarize and not any(segment.speaker is not None for segment in segments):
|
|
logger.warning("No speaker labels in the response; the model may not support diarization")
|
|
if segments:
|
|
return segments, True
|
|
text = str(data.get("text", "")).strip()
|
|
return ([Segment(start=0.0, end=duration, text=text)] if text else []), False
|
|
|
|
|
|
def turns_from_words(words: list[dict]) -> list[Segment]:
|
|
"""Group consecutive words with the same `speaker` into one Segment per speaker turn.
|
|
|
|
Words without a speaker and `audio_event` entries (e.g. laughter) are skipped.
|
|
"""
|
|
turns: list[Segment] = []
|
|
for word in words:
|
|
text = str(word.get("word", "")).strip()
|
|
speaker = word.get("speaker")
|
|
if not text or speaker is None or word.get("type") == "audio_event":
|
|
continue
|
|
if turns and turns[-1].speaker == speaker:
|
|
turns[-1].end = float(word["end"])
|
|
turns[-1].text += f" {text}"
|
|
else:
|
|
turns.append(Segment(start=float(word["start"]), end=float(word["end"]), text=text, speaker=speaker))
|
|
for turn in turns:
|
|
# Words are joined with spaces, so punctuation tokens end up detached ("Hi ,"); glue them back:
|
|
# \s+ = the stray whitespace, ([,.!?;:…]) = the punctuation mark, kept via \1.
|
|
turn.text = re.sub(r"\s+([,.!?;:…])", r"\1", turn.text)
|
|
return turns
|
|
|
|
|
|
def group_dialogue(cues: list[Segment]) -> list[tuple[float, float, str]]:
|
|
"""Turn segments into SRT cues, pairing adjacent turns by different speakers into '- a' / '- b' cues.
|
|
|
|
A pair is formed only if each line is at most DIALOGUE_LINE_LENGTH characters, the gap between them is at
|
|
most MAX_DIALOGUE_GAP seconds and the pair spans at most MAX_DIALOGUE_DURATION seconds.
|
|
Everything else becomes its own cue. Returns (start, end, text) tuples.
|
|
"""
|
|
grouped = []
|
|
index = 0
|
|
while index < len(cues):
|
|
cue = cues[index]
|
|
following = cues[index + 1] if index + 1 < len(cues) else None
|
|
if (
|
|
cue.speaker is not None
|
|
and following is not None
|
|
and following.speaker not in (None, cue.speaker)
|
|
and max(len(cue.text), len(following.text)) <= DIALOGUE_LINE_LENGTH
|
|
and following.start - cue.end <= MAX_DIALOGUE_GAP
|
|
and following.end - cue.start <= MAX_DIALOGUE_DURATION
|
|
):
|
|
grouped.append((cue.start, following.end, f"- {cue.text}\n- {following.text}"))
|
|
index += 2
|
|
else:
|
|
grouped.append((cue.start, cue.end, cue.text))
|
|
index += 1
|
|
return grouped
|
|
|
|
|
|
def split_sentences(text: str) -> list[str]:
|
|
"""Split text into sentences at SENTENCE_END matches.
|
|
|
|
A boundary is kept only if the next character starts a sentence (uppercase, digit or opening quote/bracket)
|
|
and the word before a period is not a known abbreviation or a single-letter initial.
|
|
Decimals such as "3.5" never match because no whitespace follows the period.
|
|
"""
|
|
sentences = []
|
|
start = 0
|
|
for match in SENTENCE_END.finditer(text):
|
|
following = text[match.end() : match.end() + 1]
|
|
if not following or not (following.isupper() or following.isdigit() or following in "\"'“‘(["):
|
|
continue
|
|
punctuation = match.group().strip()
|
|
# (\S+)$ = the last whitespace-delimited word before the punctuation, e.g. "Dr" in "... Dr."
|
|
word = re.search(r"(\S+)$", text[start : match.start()])
|
|
token = word.group(1).lower().lstrip("\"'“‘([") if word else ""
|
|
if punctuation.startswith(".") and (token in ABBREVIATIONS or (len(token) == 1 and token.isalpha())):
|
|
continue
|
|
sentences.append(text[start : match.start() + len(match.group().rstrip())].strip())
|
|
start = match.end()
|
|
sentences.append(text[start:].strip())
|
|
return [sentence for sentence in sentences if sentence]
|
|
|
|
|
|
def split_cue(start: float, end: float, text: str) -> list[tuple[float, float, str]]:
|
|
"""Split a cue into one cue per sentence, dividing [start, end] in proportion to character length.
|
|
|
|
Pieces shorter than MIN_SPLIT_DURATION are merged into their neighbour. The pieces are contiguous.
|
|
Returns (start, end, text) tuples; a single-sentence cue is returned unchanged.
|
|
"""
|
|
sentences = split_sentences(text=text)
|
|
if len(sentences) < 2:
|
|
return [(start, end, text)]
|
|
total = sum(len(sentence) for sentence in sentences)
|
|
pieces: list[list] = []
|
|
cursor = start
|
|
for sentence in sentences:
|
|
piece_end = cursor + (end - start) * len(sentence) / total
|
|
pieces.append([cursor, piece_end, sentence])
|
|
cursor = piece_end
|
|
pieces[-1][1] = end
|
|
merged: list[list] = []
|
|
for piece in pieces:
|
|
if merged and (piece[1] - piece[0] < MIN_SPLIT_DURATION or merged[-1][1] - merged[-1][0] < MIN_SPLIT_DURATION):
|
|
merged[-1][1] = piece[1]
|
|
merged[-1][2] += f" {piece[2]}"
|
|
else:
|
|
merged.append(piece)
|
|
return [(piece_start, piece_end, piece_text) for piece_start, piece_end, piece_text in merged]
|
|
|
|
|
|
def wrap_text(text: str) -> str:
|
|
"""Wrap text into balanced lines of at most LINE_LENGTH characters (e.g. 60 chars -> 2 lines of 30)."""
|
|
if len(text) <= LINE_LENGTH:
|
|
return text
|
|
lines_needed = -(-len(text) // LINE_LENGTH) # ceiling division
|
|
# widen the target width until the words fit in the needed number of lines
|
|
for width in range(-(-len(text) // lines_needed), LINE_LENGTH + 1):
|
|
lines = textwrap.wrap(text, width=width)
|
|
if len(lines) <= lines_needed:
|
|
break
|
|
return "\n".join(lines)
|
|
|
|
|
|
def split_long(start: float, end: float, text: str) -> list[tuple[float, float, str]]:
|
|
"""Split a cue longer than MAX_LINES lines into several cues, dividing [start, end] by character length.
|
|
|
|
Splits near the middle, preferring a position right after a comma, semicolon, colon or sentence end
|
|
(within the middle half of the text), else any space. Repeats until every piece fits.
|
|
"""
|
|
if len(text) <= LINE_LENGTH * MAX_LINES:
|
|
return [(start, end, text)]
|
|
middle = len(text) / 2
|
|
spaces = [match.start() for match in re.finditer(" ", text)]
|
|
if not spaces:
|
|
return [(start, end, text)]
|
|
# a space preceded by , ; : . ! ? is a clause break
|
|
clause_breaks = [position for position in spaces if text[position - 1] in ",;:.!?" and abs(position - middle) <= len(text) / 4]
|
|
position = min(clause_breaks or spaces, key=lambda candidate: abs(candidate - middle))
|
|
left, right = text[:position].strip(), text[position:].strip()
|
|
boundary = start + (end - start) * len(left) / (len(left) + len(right))
|
|
return split_long(start=start, end=boundary, text=left) + split_long(start=boundary, end=end, text=right)
|
|
|
|
|
|
def format_cues(cues: list[tuple[float, float, str]]) -> list[tuple[float, float, str]]:
|
|
"""Split single-speaker cues that are too long and wrap their text into lines; dialogue cues pass through."""
|
|
formatted = []
|
|
for start, end, text in cues:
|
|
if text.startswith("- "):
|
|
formatted.append((start, end, text))
|
|
else:
|
|
formatted.extend(
|
|
(piece_start, piece_end, wrap_text(text=piece_text)) for piece_start, piece_end, piece_text in split_long(start=start, end=end, text=text)
|
|
)
|
|
return formatted
|
|
|
|
|
|
def format_timestamp(seconds: float) -> str:
|
|
"""Format seconds as an SRT timestamp, HH:MM:SS,mmm."""
|
|
total_milliseconds = 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 build_srt(cues: list[tuple[float, float, str]]) -> str:
|
|
"""Render (start, end, text) cues as SRT, keeping them non-overlapping with a minimum length of 0.3 s."""
|
|
blocks = []
|
|
previous_end = 0.0
|
|
for index, (start, end, text) in enumerate(cues, start=1):
|
|
start = max(start, previous_end)
|
|
end = max(end, start + 0.3)
|
|
previous_end = end
|
|
blocks.append(f"{index}\n{format_timestamp(start)} --> {format_timestamp(end)}\n{text}\n")
|
|
return "\n".join(blocks)
|
|
|
|
|
|
def main() -> None:
|
|
"""Extract audio, transcribe it chunk by chunk, rescale timestamps by the speed-up and write the SRT."""
|
|
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
|
|
args = parse_args()
|
|
if not args.api_key:
|
|
raise SystemExit("Provide --api-key or set OPENROUTER_API_KEY")
|
|
if not args.video.is_file():
|
|
raise SystemExit(f"Input not found: {args.video}")
|
|
if not 1.0 <= args.speed_up <= 2.0:
|
|
raise SystemExit("--speed-up must be between 1.0 and 2.0")
|
|
|
|
output = args.output or args.video.with_suffix(".srt")
|
|
cues: list[Segment] = []
|
|
|
|
with (
|
|
tempfile.TemporaryDirectory() as temp_directory,
|
|
httpx.Client(
|
|
base_url="https://openrouter.ai/api/v1",
|
|
headers={"Authorization": f"Bearer {args.api_key}"},
|
|
timeout=300,
|
|
) as client,
|
|
):
|
|
chunks = run_ffmpeg_extract(
|
|
video=args.video,
|
|
directory=Path(temp_directory),
|
|
chunk_seconds=int(args.chunk_minutes * 60),
|
|
speed_up=args.speed_up,
|
|
)
|
|
logger.info("Extracted %d audio chunk(s)", len(chunks))
|
|
offset = 0.0
|
|
for index, chunk in enumerate(chunks, start=1):
|
|
duration = get_duration(path=chunk)
|
|
if duration < MIN_CHUNK_DURATION:
|
|
logger.info("Skipping empty chunk %d/%d", index, len(chunks))
|
|
continue
|
|
logger.info("Transcribing chunk %d/%d with %s", index, len(chunks), args.model)
|
|
segments, has_timestamps = transcribe_chunk(
|
|
client=client,
|
|
path=chunk,
|
|
duration=duration,
|
|
model=args.model,
|
|
language=args.language,
|
|
diarize=args.diarize,
|
|
)
|
|
for segment in segments:
|
|
start = segment.start * args.speed_up + offset
|
|
end = segment.end * args.speed_up + offset
|
|
if args.speed_up > 1.0 or args.diarize or not has_timestamps:
|
|
pieces = split_cue(start=start, end=end, text=segment.text)
|
|
else:
|
|
pieces = [(start, end, segment.text)]
|
|
cues.extend(Segment(start=piece[0], end=piece[1], text=piece[2], speaker=segment.speaker) for piece in pieces)
|
|
offset += duration * args.speed_up
|
|
|
|
if not cues:
|
|
raise SystemExit("Transcription returned no text")
|
|
srt_cues = format_cues(cues=group_dialogue(cues=cues))
|
|
output.write_text(build_srt(cues=srt_cues), encoding="utf-8")
|
|
logger.info("Wrote %d cues to %s", len(srt_cues), output)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|