"""Align a subtitle against a few audio windows of a stream.
Pure functions, no I/O: addon.py fetches the windows and runs VAD on them
(ffsubsync --serialize-speech), this decides what they say. See
test_window_align.py for why ffsubsync's own search is not used here.
Model (ffsubsync's, so results read the same as the library sync's):
corrected time = scale * t + offset
"""
import re
from typing import Optional
import numpy as np
SAMPLE_RATE = 100 # ffsubsync's speech arrays are 100 Hz
# 1.0 first: on a tie it wins, and an exact 1.0 is what "in sync" compares.
SCALES = (1.0, 24 / 23.976, 23.976 / 24, 25 / 23.976, 23.976 / 25, 25 / 24, 24 / 25)
MAX_LAG = 60.0 # seconds searched either side
MIN_SCORE = 0.15 # measured: matching windows 0.17-0.60,
# wrong ones -0.16-0.09 (Hope, NOC)
AGREE = 0.5 # seconds; NOC en windows spread 0.95
# only where the audio had a gap
SPEECH_FRAC = (0.05, 0.9) # outside this the VAD heard music/nothing
_TS = r"(\d+):(\d{2}):(\d{2})[,.](\d{1,3})"
_CUE = re.compile(_TS + r"(\s*-->\s*)" + _TS)
def _secs(h, m, s, ms) -> float:
return int(h) * 3600 + int(m) * 60 + int(s) + int(ms.ljust(3, "0")) / 1000
def srt_cues(text: str) -> list:
"""[(start, end)] in seconds, in file order."""
return [(_secs(*m.group(1, 2, 3, 4)), _secs(*m.group(6, 7, 8, 9)))
for m in _CUE.finditer(text)]
def _signal(cues, scale: float, start: float, length: int) -> np.ndarray:
sig = np.zeros(length)
for a, b in cues:
i = int(round((a * scale - start) * SAMPLE_RATE))
j = int(round((b * scale - start) * SAMPLE_RATE))
i, j = max(i, 0), min(j, length)
if i < j:
sig[i:j] = 1.0
return sig
def _lag(cues, scale: float, pos: float, speech: np.ndarray) -> tuple:
"""(offset, score) that best lines the scaled cues up with one window.
Score is the ±1 agreement per sample, so -1..1 whatever the length."""
pad = int(MAX_LAG * SAMPLE_RATE)
n = len(speech)
w = 2 * speech - 1
s = 2 * _signal(cues, scale, pos - MAX_LAG, n + 2 * pad) - 1
size = 1 << int(np.ceil(np.log2(len(s) + n)))
corr = np.fft.irfft(np.fft.rfft(s, size) * np.conj(np.fft.rfft(w, size)), size)
valid = corr[:2 * pad + 1]
k = int(np.argmax(valid))
return (pad - k) / SAMPLE_RATE, float(valid[k]) / n
def align(cues, windows) -> Optional[dict]:
"""{"offset", "scale", "windows", "score"} or None when the windows do not
agree on an answer. `windows` is [(position_seconds, speech_array)]."""
if not cues:
return None
usable = []
for pos, speech in windows:
speech = (np.asarray(speech, dtype=float) >= 1.0).astype(float)
if len(speech) and SPEECH_FRAC[0] < speech.mean() < SPEECH_FRAC[1]:
usable.append((pos, speech))
best = None
for scale in SCALES:
good = [r for r in (_lag(cues, scale, p, w) for p, w in usable)
if r[1] >= MIN_SCORE]
if len(good) < 2:
continue
cluster = max(([g for g in good if abs(g[0] - o) <= AGREE] for o, _ in good),
key=lambda c: (len(c), sum(x[1] for x in c)))
if len(cluster) < 2 or 2 * len(cluster) < len(good):
continue
cand = {"offset": round(float(np.median([c[0] for c in cluster])), 2),
"scale": scale, "windows": len(cluster),
"score": round(sum(c[1] for c in cluster) / len(cluster), 3)}
if best is None or (cand["windows"], cand["score"]) > (best["windows"], best["score"]):
best = cand
return best
def _fmt(t: float) -> str:
ms = max(0, int(round(t * 1000)))
h, ms = divmod(ms, 3_600_000)
m, ms = divmod(ms, 60_000)
s, ms = divmod(ms, 1000)
return f"{h:02d}:{m:02d}:{s:02d},{ms:03d}"
def apply(text: str, scale: float, offset: float) -> str:
"""The SRT with every timestamp moved to scale * t + offset (>= 0)."""
def one(m):
a = _secs(*m.group(1, 2, 3, 4)) * scale + offset
b = _secs(*m.group(6, 7, 8, 9)) * scale + offset
return _fmt(a) + m.group(5) + _fmt(b)
return _CUE.sub(one, text)