feat: Add subtitle timing and segmentation script
This commit is contained in:
abdus committed
1 parent
be9670b653
commit
f867799876
1 file changed
+364
Executable
+364
@@ -0,0 +1,364 @@
|
||||
#!/usr/bin/env python3
|
||||
import argparse
|
||||
import logging
|
||||
import re
|
||||
import textwrap
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
# 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,
|
||||
)
|
||||
# Matches an SRT time line such as "00:01:02,345 --> 00:01:04,000" and captures the 8 time fields.
|
||||
TIMELINE_PATTERN = re.compile(
|
||||
r"""
|
||||
(\d+):(\d{2}):(\d{2})[,.](\d{3}) # start: hours, minutes, seconds, milliseconds (comma or dot)
|
||||
\s*-->\s* # arrow, with optional spaces around it
|
||||
(\d+):(\d{2}):(\d{2})[,.](\d{3}) # end, same layout
|
||||
""",
|
||||
re.VERBOSE,
|
||||
)
|
||||
# Matches a whole cue that is only a sound description or music, e.g. "[Music]", "(laughing)", "♪ la la ♪".
|
||||
NONSPEECH_PATTERN = re.compile(
|
||||
r"""
|
||||
^(
|
||||
\[[^\]]*\] # [ ... ]
|
||||
| \([^)]*\) # ( ... )
|
||||
| ♪[^♪]*♪? # ♪ ... ♪, closing note optional
|
||||
| ♫[^♫]*♫? # ♫ ... ♫, closing note optional
|
||||
)$
|
||||
""",
|
||||
re.VERBOSE,
|
||||
)
|
||||
# Matches inline markup to strip: HTML tags like <i> and </i>, and ASS/SSA overrides like {\an8}.
|
||||
MARKUP_PATTERN = re.compile(
|
||||
r"""
|
||||
<[^>]+> # < ... >
|
||||
| \{[^}]*\} # { ... }
|
||||
""",
|
||||
re.VERBOSE,
|
||||
)
|
||||
MIN_CUE_DURATION = 0.2
|
||||
|
||||
logger = logging.getLogger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class Cue:
|
||||
"""A subtitle cue or piece of one; times in seconds. `dialogue` marks cues holding several '- ' speaker lines."""
|
||||
|
||||
start: float
|
||||
end: float
|
||||
text: str
|
||||
dialogue: bool = False
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
"""Parse command-line arguments."""
|
||||
parser = argparse.ArgumentParser(
|
||||
description="Re-time and re-segment an SRT file mechanically: merge by sentence, split long cues, wrap lines, enforce reading speed.",
|
||||
formatter_class=argparse.ArgumentDefaultsHelpFormatter,
|
||||
)
|
||||
parser.add_argument("subtitle", type=Path, help="Input SRT file")
|
||||
parser.add_argument("--output", "-o", type=Path, help="Output SRT path (default: <name>.rewritten.srt)")
|
||||
parser.add_argument("--max-line-length", type=int, default=42, help="Max characters per line")
|
||||
parser.add_argument("--max-lines", type=int, default=2, help="Max lines per cue")
|
||||
parser.add_argument("--max-cps", type=float, default=17.0, help="Max characters per second before a cue's end is extended")
|
||||
parser.add_argument("--min-duration", type=float, default=1.0, help="Cues shorter than this (seconds) are merged into a neighbour")
|
||||
parser.add_argument("--max-duration", type=float, default=7.0, help="Cues longer than this (seconds) are split")
|
||||
parser.add_argument("--min-gap", type=float, default=0.04, help="Minimum gap between cues in seconds")
|
||||
parser.add_argument("--drop-nonspeech", action="store_true", help="Remove cues that are only [..], (..) or music notes")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def parse_timestamp(hours: str, minutes: str, seconds: str, milliseconds: str) -> float:
|
||||
"""Convert the four captured time fields of TIMELINE_PATTERN to seconds."""
|
||||
return int(hours) * 3600 + int(minutes) * 60 + int(seconds) + int(milliseconds) / 1000
|
||||
|
||||
|
||||
def parse_srt(content: str) -> list[Cue]:
|
||||
"""Parse SRT text into cues sorted by start time.
|
||||
|
||||
Tolerates CRLF line endings and a missing index line. Text lines are flattened into one line and markup is
|
||||
stripped. A cue with two or more lines starting with '- ' is flagged as dialogue. Empty cues are dropped.
|
||||
"""
|
||||
cues = []
|
||||
# \n\s*\n = a blank line (possibly containing spaces) separates SRT blocks
|
||||
for block in re.split(r"\n\s*\n", content.replace("\r\n", "\n").replace("\r", "\n").strip()):
|
||||
lines = block.split("\n")
|
||||
match = next((found for line in lines[:3] if (found := TIMELINE_PATTERN.search(line))), None)
|
||||
if match is None:
|
||||
raise ValueError(f"Invalid SRT block: {block[:80]!r}")
|
||||
position = next(position for position, line in enumerate(lines[:3]) if TIMELINE_PATTERN.search(line))
|
||||
body = [line.strip() for line in lines[position + 1 :] if line.strip()]
|
||||
# \s+ -> " " collapses runs of whitespace (including the gaps left by removed markup) to one space
|
||||
text = re.sub(r"\s+", " ", MARKUP_PATTERN.sub("", " ".join(body))).strip()
|
||||
if text:
|
||||
cues.append(
|
||||
Cue(
|
||||
start=parse_timestamp(*match.group(1, 2, 3, 4)),
|
||||
end=parse_timestamp(*match.group(5, 6, 7, 8)),
|
||||
text=text,
|
||||
dialogue=sum(line.startswith("- ") for line in body) >= 2,
|
||||
)
|
||||
)
|
||||
if not cues:
|
||||
raise ValueError("SRT contains no cues")
|
||||
cues.sort(key=lambda cue: cue.start)
|
||||
return cues
|
||||
|
||||
|
||||
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 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, opening quote/bracket or a
|
||||
'- ' speaker dash) 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]
|
||||
starts_sentence = following.isupper() or following.isdigit() or following in "\"'“‘([" or text.startswith("- ", match.end())
|
||||
if not following or not starts_sentence:
|
||||
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 wrap_text(text: str, width: int, max_lines: int) -> list[str] | None:
|
||||
"""Wrap text to at most `max_lines` lines of at most `width` characters, preferring balanced lines.
|
||||
|
||||
Text of the form "- a - b" is split into one line per speaker when each fits. Returns None if no wrapping fits.
|
||||
"""
|
||||
if len(text) <= width:
|
||||
return [text]
|
||||
if text.startswith("- ") and " - " in text:
|
||||
# \s+(?=- ) = whitespace followed by a dash-space (lookahead, not consumed): splits before each speaker dash
|
||||
speakers = re.split(r"\s+(?=- )", text)
|
||||
if len(speakers) <= max_lines and all(len(line) <= width for line in speakers):
|
||||
return speakers
|
||||
for count in range(2, max_lines + 1):
|
||||
lines = textwrap.wrap(text, width=-(-len(text) // count), break_long_words=False, break_on_hyphens=False)
|
||||
if len(lines) <= max_lines and all(len(line) <= width for line in lines):
|
||||
return lines
|
||||
lines = textwrap.wrap(text, width=width, break_long_words=False, break_on_hyphens=False)
|
||||
return lines if len(lines) <= max_lines and all(len(line) <= width for line in lines) else None
|
||||
|
||||
|
||||
def wrap_dialogue(text: str, width: int) -> list[str]:
|
||||
"""Wrap a dialogue cue, one or more lines per speaker. The number of lines is not limited."""
|
||||
lines = []
|
||||
# \s+(?=- ) = whitespace followed by a dash-space (lookahead, not consumed): splits before each speaker dash
|
||||
for speaker in re.split(r"\s+(?=- )", text):
|
||||
lines.extend(textwrap.wrap(speaker, width=width, break_long_words=False, break_on_hyphens=False))
|
||||
return lines
|
||||
|
||||
|
||||
def time_at(spans: list[tuple[int, Cue]], offset: int, is_end: bool) -> float:
|
||||
"""Map a character offset in the joined text to a time, interpolating linearly inside the cue that holds it.
|
||||
|
||||
`spans` are (offset of cue text in the joined text, cue). For an end offset the cue boundary belongs to the
|
||||
earlier cue, for a start offset to the later one. Offsets in the joining gap snap to the nearest cue edge.
|
||||
"""
|
||||
for span_offset, cue in spans:
|
||||
length = len(cue.text)
|
||||
inside = span_offset < offset <= span_offset + length if is_end else span_offset <= offset < span_offset + length
|
||||
if inside:
|
||||
return cue.start + (cue.end - cue.start) * (offset - span_offset) / length
|
||||
span_offset, cue = min(spans, key=lambda span: min(abs(offset - span[0]), abs(offset - span[0] - len(span[1].text))))
|
||||
return cue.start if offset <= span_offset else cue.end
|
||||
|
||||
|
||||
def sentences_with_times(cues: list[Cue]) -> list[Cue]:
|
||||
"""Join the cues' text, split it into sentences and give each sentence a start and end time.
|
||||
|
||||
A sentence spanning several cues starts where its first character is spoken and ends where its last one is.
|
||||
"""
|
||||
joined = " ".join(cue.text for cue in cues)
|
||||
spans = []
|
||||
offset = 0
|
||||
for cue in cues:
|
||||
spans.append((offset, cue))
|
||||
offset += len(cue.text) + 1
|
||||
pieces = []
|
||||
cursor = 0
|
||||
for sentence in split_sentences(text=joined):
|
||||
position = joined.index(sentence, cursor)
|
||||
cursor = position + len(sentence)
|
||||
pieces.append(
|
||||
Cue(
|
||||
start=time_at(spans=spans, offset=position, is_end=False),
|
||||
end=time_at(spans=spans, offset=cursor, is_end=True),
|
||||
text=sentence,
|
||||
)
|
||||
)
|
||||
return pieces
|
||||
|
||||
|
||||
def fits(text: str, duration: float, args: argparse.Namespace) -> bool:
|
||||
"""Whether one cue can hold `text`: within --max-duration and wrappable into --max-lines lines."""
|
||||
return duration <= args.max_duration and wrap_text(text=text, width=args.max_line_length, max_lines=args.max_lines) is not None
|
||||
|
||||
|
||||
def merge_short(pieces: list[Cue], args: argparse.Namespace) -> list[Cue]:
|
||||
"""Merge a piece into the previous one when either is shorter than --min-duration and the result still fits."""
|
||||
merged: list[Cue] = []
|
||||
for piece in pieces:
|
||||
if merged:
|
||||
previous = merged[-1]
|
||||
short = piece.end - piece.start < args.min_duration or previous.end - previous.start < args.min_duration
|
||||
text = f"{previous.text} {piece.text}"
|
||||
if short and fits(text=text, duration=piece.end - previous.start, args=args):
|
||||
previous.end = piece.end
|
||||
previous.text = text
|
||||
continue
|
||||
merged.append(piece)
|
||||
return merged
|
||||
|
||||
|
||||
def find_split(text: str) -> int | None:
|
||||
"""Pick the index at which to cut `text` in two, or None if it is a single word.
|
||||
|
||||
Prefers the start of the word after a comma, semicolon or colon if it is within 25% of the length from the
|
||||
middle; otherwise the word start nearest the middle.
|
||||
"""
|
||||
middle = len(text) / 2
|
||||
# \s+ = any whitespace run; match.end() is where the next word starts (interior only, not at either edge)
|
||||
word_starts = [match.end() for match in re.finditer(r"\s+", text) if 0 < match.start() and match.end() < len(text)]
|
||||
if not word_starts:
|
||||
return None
|
||||
# [,;:]\s+ = a clause mark followed by whitespace; match.end() is the start of the next clause
|
||||
clause_starts = [match.end() for match in re.finditer(r"[,;:]\s+", text) if match.end() < len(text)]
|
||||
if clause_starts:
|
||||
best = min(clause_starts, key=lambda index: abs(index - middle))
|
||||
if abs(best - middle) <= len(text) * 0.25:
|
||||
return best
|
||||
return min(word_starts, key=lambda index: abs(index - middle))
|
||||
|
||||
|
||||
def split_piece(piece: Cue, args: argparse.Namespace) -> list[Cue]:
|
||||
"""Recursively split a piece until every part fits, dividing its time in proportion to character length."""
|
||||
if fits(text=piece.text, duration=piece.end - piece.start, args=args):
|
||||
return [piece]
|
||||
index = find_split(text=piece.text)
|
||||
if index is None:
|
||||
return [piece]
|
||||
left, right = piece.text[:index].strip(), piece.text[index:].strip()
|
||||
middle = piece.start + (piece.end - piece.start) * len(left) / (len(left) + len(right))
|
||||
return split_piece(piece=Cue(start=piece.start, end=middle, text=left), args=args) + split_piece(
|
||||
piece=Cue(start=middle, end=piece.end, text=right), args=args
|
||||
)
|
||||
|
||||
|
||||
def retime(pieces: list[Cue], args: argparse.Namespace) -> list[int]:
|
||||
"""Fix timing in place: remove overlaps, keep --min-gap, and extend a cue's end to meet --max-cps.
|
||||
|
||||
An end is only extended into free time before the next cue and never beyond --max-duration.
|
||||
Returns the 1-based numbers of cues still above --max-cps.
|
||||
"""
|
||||
over_budget = []
|
||||
for index, piece in enumerate(pieces):
|
||||
if index and piece.start < pieces[index - 1].end:
|
||||
piece.start = pieces[index - 1].end
|
||||
limit = pieces[index + 1].start - args.min_gap if index + 1 < len(pieces) else float("inf")
|
||||
wanted = piece.start + len(piece.text) / args.max_cps
|
||||
ceiling = min(piece.start + args.max_duration, limit)
|
||||
if wanted > piece.end:
|
||||
piece.end = max(piece.end, min(wanted, ceiling))
|
||||
if piece.end > limit:
|
||||
piece.end = max(limit, piece.start + MIN_CUE_DURATION)
|
||||
if piece.end <= piece.start:
|
||||
piece.end = piece.start + MIN_CUE_DURATION
|
||||
if len(piece.text) / (piece.end - piece.start) > args.max_cps + 0.01:
|
||||
over_budget.append(index + 1)
|
||||
return over_budget
|
||||
|
||||
|
||||
def rewrite(cues: list[Cue], args: argparse.Namespace) -> tuple[list[Cue], list[int]]:
|
||||
"""Run the full pipeline. Dialogue and non-speech cues act as barriers and pass through unchanged;
|
||||
the speech between them is merged into sentences, short pieces are merged, long ones split, then retimed.
|
||||
|
||||
Returns (pieces, numbers of cues still above --max-cps).
|
||||
"""
|
||||
pieces: list[Cue] = []
|
||||
run: list[Cue] = []
|
||||
|
||||
def flush() -> None:
|
||||
if run:
|
||||
pieces.extend(
|
||||
part
|
||||
for piece in merge_short(pieces=sentences_with_times(cues=run), args=args)
|
||||
for part in split_piece(piece=piece, args=args)
|
||||
)
|
||||
run.clear()
|
||||
|
||||
for cue in cues:
|
||||
if cue.dialogue or NONSPEECH_PATTERN.match(cue.text):
|
||||
flush()
|
||||
if cue.dialogue or not args.drop_nonspeech:
|
||||
pieces.append(Cue(start=cue.start, end=cue.end, text=cue.text, dialogue=cue.dialogue))
|
||||
else:
|
||||
run.append(cue)
|
||||
flush()
|
||||
return pieces, retime(pieces=pieces, args=args)
|
||||
|
||||
|
||||
def build_srt(pieces: list[Cue], args: argparse.Namespace) -> str:
|
||||
"""Render the pieces as SRT, wrapping each to --max-line-length (dialogue cues one speaker per line)."""
|
||||
blocks = []
|
||||
for index, piece in enumerate(pieces, start=1):
|
||||
if piece.dialogue:
|
||||
lines = wrap_dialogue(text=piece.text, width=args.max_line_length)
|
||||
else:
|
||||
lines = wrap_text(text=piece.text, width=args.max_line_length, max_lines=args.max_lines) or [piece.text]
|
||||
blocks.append(f"{index}\n{format_timestamp(piece.start)} --> {format_timestamp(piece.end)}\n" + "\n".join(lines) + "\n")
|
||||
return "\n".join(blocks)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Read the SRT, rewrite it and write the result next to it, refusing to overwrite the input."""
|
||||
logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s: %(message)s")
|
||||
args = parse_args()
|
||||
if min(args.max_line_length, args.max_lines) < 1 or args.max_cps <= 0:
|
||||
raise SystemExit("--max-line-length, --max-lines and --max-cps must be positive")
|
||||
output = args.output or args.subtitle.with_name(f"{args.subtitle.stem}.rewritten.srt")
|
||||
if output.resolve() == args.subtitle.resolve():
|
||||
raise SystemExit("Refusing to overwrite the input file")
|
||||
|
||||
cues = parse_srt(content=args.subtitle.read_text(encoding="utf-8-sig"))
|
||||
pieces, over_budget = rewrite(cues=cues, args=args)
|
||||
output.write_text(build_srt(pieces=pieces, args=args), encoding="utf-8")
|
||||
logger.info("Wrote %d cues (from %d) to %s", len(pieces), len(cues), output)
|
||||
if over_budget:
|
||||
logger.warning("%d cue(s) still above %.1f chars/s (text is not shortened): %s", len(over_budget), args.max_cps, over_budget[:20])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in new issue
Block a user