memby/server/internal/subsync/subtitles.go

// Package subsync fixes the timing of a subtitle by aligning it to one that is already
// right.
//
// The idea is borrowed from ffsubsync and alass, which AutoSubSync drives — but not the
// code, which is Python and Rust and needs FFmpeg to read a film's audio. This gateway's
// image is distroless with no ffmpeg and no shell, so the audio path is closed to it. What
// is open is the cheaper half of the same idea: both of those tools can align against a
// *reference subtitle* instead of audio, and a household's copy of a film usually carries
// a track that is already correct. Reduce both tracks to "somebody is speaking / nobody
// is" on a fixed time grid and the offset is the shift where the two agree most.
//
// Nothing here touches the network, Emby or the database, which is what lets the whole
// rule be tested against real subtitle text.
package subsync

import (
	"fmt"
	"sort"
	"strings"
	"time"
)

// Cue is one subtitle line: when it appears, when it goes, and what it says.
//
// The text is carried through untouched. This package changes *when* a line is shown and
// never what it says — a resynchronised subtitle that had also been reflowed or re-escaped
// would be impossible to tell from a corrupted one.
type Cue struct {
	Start time.Duration
	End   time.Duration
	Text  string
}

// Duration is where the last cue ends, which stands in for the runtime the track was
// written against.
func Duration(cues []Cue) time.Duration {
	var last time.Duration
	for _, cue := range cues {
		if cue.End > last {
			last = cue.End
		}
	}
	return last
}

// Shift moves every cue by the correction Align found.
//
// Scale is applied before the offset, in that order, because that is the order the damage
// happened in: a subtitle written for a 25fps transfer and played against a 23.976 one
// runs progressively further out, and whatever fixed offset remains sits on top of the
// stretch. A cue dragged before zero is clamped rather than dropped — it belongs to a line
// somebody is about to hear, and a track that silently lost its opening would look like a
// worse fault than the one being fixed.
func Shift(cues []Cue, offset time.Duration, scale float64) []Cue {
	if scale <= 0 {
		scale = 1
	}
	out := make([]Cue, 0, len(cues))
	for _, cue := range cues {
		start := scaleDuration(cue.Start, scale) + offset
		end := scaleDuration(cue.End, scale) + offset
		if end < 0 {
			// The whole cue was pushed off the front. Keep it at zero rather than
			// discarding it: losing a line is worse than showing one early.
			start, end = 0, 0
		} else if start < 0 {
			start = 0
		}
		out = append(out, Cue{Start: start, End: end, Text: cue.Text})
	}
	sort.SliceStable(out, func(i, j int) bool { return out[i].Start < out[j].Start })
	return out
}

func scaleDuration(d time.Duration, scale float64) time.Duration {
	return time.Duration(float64(d) * scale)
}

// FrameRateScales are the stretch factors tried alongside a plain offset.
//
// They are the ratios between the frame rates films are actually delivered at, because
// the classic broken subtitle is not merely late — it is a track authored against one
// transfer and played against another, so it drifts, and no single offset can fix it. The
// list is deliberately short: every extra candidate is another chance for a wrong answer
// to win by luck, and these cover the transfers that exist.
var FrameRateScales = []float64{
	1,
	23.976 / 24, 24 / 23.976,
	25.0 / 24, 24.0 / 25,
	25 / 23.976, 23.976 / 25,
	30 / 29.97, 29.97 / 30,
}

// Options tunes the search. The zero value is not usable; use DefaultOptions.
type Options struct {
	// Bin is the width of one cell of the time grid. Fine enough that a correction is
	// worth making at all, coarse enough that two tracks written by different people
	// still land in the same cells.
	Bin time.Duration
	// MaxShift bounds the search either way. A subtitle for a different release can be a
	// minute out; beyond that the two are far more likely to be different cuts of the
	// film, where a confident answer would be a wrong one.
	MaxShift time.Duration
	// MinCues is how much evidence is needed on each side before an answer is offered.
	MinCues int
	// MinScore is the share of the shorter track's speech that must line up.
	MinScore float64
	// MinMargin is how far the winning shift must beat the best rival outside its own
	// peak. This is what separates a real alignment from a track that correlates weakly
	// with everything, which is what a wrong reference looks like.
	MinMargin float64
	// Scales are the stretch candidates tried. Empty means offset only.
	Scales []float64
}

// DefaultOptions is the tuning the gateway uses.
func DefaultOptions() Options {
	return Options{
		Bin:       50 * time.Millisecond,
		MaxShift:  60 * time.Second,
		MinCues:   20,
		MinScore:  0.45,
		MinMargin: 0.10,
		Scales:    FrameRateScales,
	}
}

// Result describes a correction.
type Result struct {
	// Offset is what to add to every cue, after Scale.
	Offset time.Duration
	// Scale is the stretch applied first. 1 means the timing was merely displaced.
	Scale float64
	// Score is the share of the shorter track's speech that lines up once corrected —
	// how much the two agree, not how confident we are.
	Score float64
	// Margin is how far ahead of the best rival alignment this one finished. A high
	// score with no margin is a track that matches everything, which matches nothing.
	Margin float64
}

// Correction reports whether this result actually changes anything a viewer would see.
func (r Result) Correction() bool {
	return r.Offset.Abs() >= 100*time.Millisecond || r.Scale != 1
}

// String renders the correction the way the console and the log want it.
func (r Result) String() string {
	seconds := r.Offset.Seconds()
	out := fmt.Sprintf("%+.2fs", seconds)
	if r.Scale != 1 {
		out += fmt.Sprintf(" at %.4f×", r.Scale)
	}
	return out
}

// ErrNoAlignment is returned when the two tracks cannot be aligned with confidence.
//
// Most of this package is about producing this rather than a number. A subtitle nudged to
// the wrong place is worse than one left alone: the viewer asked for a fix, would be told
// it worked, and would have no way of knowing the file they now have is further out than
// the one they started with.
type ErrNoAlignment struct {
	// Reason is safe to show a viewer.
	Reason string
	// Score and Margin are what the search actually found, for the log.
	Score  float64
	Margin float64
}

func (e *ErrNoAlignment) Error() string {
	return fmt.Sprintf("subsync: %s (score %.2f, margin %.2f)", e.Reason, e.Score, e.Margin)
}

// Align finds the correction that brings broken onto reference.
//
// The reference is assumed correct; nothing here checks that, because nothing could. The
// caller chooses it, and choosing badly is the one failure this package cannot detect —
// which is why the margin test exists, since a wrong reference tends to match everything
// equally rather than matching one shift particularly well.
func Align(broken, reference []Cue, opts Options) (Result, error) {
	if opts.Bin <= 0 {
		opts = DefaultOptions()
	}
	if len(broken) < opts.MinCues || len(reference) < opts.MinCues {
		return Result{}, &ErrNoAlignment{
			Reason: "there are not enough lines in one of these subtitles to compare them",
		}
	}

	scales := opts.Scales
	if len(scales) == 0 {
		scales = []float64{1}
	}

	// The reference is rasterised once; the broken track is rasterised per stretch
	// candidate, which is a handful of times rather than once per shift.
	ref := rasterise(reference, opts.Bin, 1)
	maxShiftBins := int(opts.MaxShift / opts.Bin)

	best := Result{Scale: 1}
	bestBins := 0
	found := false
	// Rivals holds the best score at each shift far enough from the winner to be a
	// different answer rather than the same peak's shoulder.
	var runnerUp float64

	for _, scale := range scales {
		signal := rasterise(broken, opts.Bin, scale)
		if signal.on == 0 {
			continue
		}
		floor := min(signal.on, ref.on)
		if floor == 0 {
			continue
		}
		for shift := -maxShiftBins; shift <= maxShiftBins; shift++ {
			score := float64(signal.overlap(ref, shift)) / float64(floor)
			switch {
			case !found || score > best.Score:
				// The old winner becomes a rival only if it is a different answer.
				if found && farApart(bestBins, shift, best.Scale, scale, opts.Bin) {
					runnerUp = max(runnerUp, best.Score)
				}
				found = true
				best = Result{
					Offset: time.Duration(shift) * opts.Bin,
					Scale:  scale,
					Score:  score,
				}
				bestBins = shift
			case farApart(bestBins, shift, best.Scale, scale, opts.Bin):
				runnerUp = max(runnerUp, score)
			}
		}
	}

	if !found {
		return Result{}, &ErrNoAlignment{Reason: "these subtitles have no speech in common"}
	}
	best.Margin = best.Score - runnerUp
	if best.Score < opts.MinScore {
		return Result{}, &ErrNoAlignment{
			Reason: "these two subtitles are too different to line up — they may be for " +
				"different cuts of this title",
			Score: best.Score, Margin: best.Margin,
		}
	}
	if best.Margin < opts.MinMargin {
		return Result{}, &ErrNoAlignment{
			Reason: "no single timing fits better than the others, so the result would be " +
				"a guess",
			Score: best.Score, Margin: best.Margin,
		}
	}
	return best, nil
}

// farApart says whether two candidates are different answers rather than two samples of
// one peak. Alignment produces a broad hill either side of the true shift, so the shifts
// immediately around the winner are not rivals — treating them as rivals would collapse
// every margin to nearly zero and refuse every correct answer.
func farApart(bestBins, shift int, bestScale, scale float64, bin time.Duration) bool {
	if bestScale != scale {
		return true
	}
	guard := int(time.Second / bin)
	return abs(shift-bestBins) > guard
}

func abs(v int) int {
	if v < 0 {
		return -v
	}
	return v
}

// Reference picks which of the tracks on offer to align against.
//
// The rules are about what makes a *usable* yardstick, not what makes a good subtitle. A
// forced track carries only the lines that are foreign to the film's own audio, so it is
// mostly silence and would correlate with almost any shift; a track with very few cues is
// the same problem in a different shape. Longest wins because coverage is what the
// correlation is measured against. It returns the index so the caller can name the track
// it used, which is the one thing a viewer needs to judge the answer.
func Reference(candidates [][]Cue, exclude int, forced []bool, minCues int) (int, bool) {
	best, bestCount := -1, 0
	for i, cues := range candidates {
		if i == exclude || len(cues) < minCues {
			continue
		}
		if i < len(forced) && forced[i] {
			continue
		}
		if len(cues) > bestCount {
			best, bestCount = i, len(cues)
		}
	}
	return best, best >= 0
}

// rasterise turns cues into one bit per time bin: is anybody speaking in this cell.
//
// The text is thrown away deliberately. Two subtitles for one film are in different
// languages as often as not, so the only thing they can be compared on is *when* lines
// happen — which turns out to be plenty, because dialogue rhythm is a property of the
// film rather than of the translation.
func rasterise(cues []Cue, bin time.Duration, scale float64) *signal {
	end := time.Duration(0)
	for _, cue := range cues {
		if e := scaleDuration(cue.End, scale); e > end {
			end = e
		}
	}
	sig := newSignal(int(end/bin) + 2)
	for _, cue := range cues {
		start := scaleDuration(cue.Start, scale)
		stop := scaleDuration(cue.End, scale)
		if stop <= start {
			// A zero-length or reversed cue still says a line happened here.
			stop = start + bin
		}
		sig.set(int(start/bin), int(stop/bin))
	}
	return sig
}

// Normalise tidies a parsed track before it is compared or written back.
//
// Cues out of order, or overlapping, are ordinary in files people have edited by hand, and
// both would put speech in the wrong cell.
func Normalise(cues []Cue) []Cue {
	out := make([]Cue, 0, len(cues))
	for _, cue := range cues {
		if strings.TrimSpace(cue.Text) == "" {
			continue
		}
		if cue.End < cue.Start {
			cue.Start, cue.End = cue.End, cue.Start
		}
		out = append(out, cue)
	}
	sort.SliceStable(out, func(i, j int) bool { return out[i].Start < out[j].Start })
	return out
}