feat: Add video transcription script
This commit is contained in:
abdus committed
1 parent
f6381d4c35
commit
e59f2d4825
1 file changed
+468
Executable
+468
@@ -0,0 +1,468 @@
|
||||
#!/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()
|
||||
Reference in new issue
Block a user