old-books/_pipeline/assemble.py

"""Assembly with ffmpeg: wav pieces + chimes + pauses -> one MP3 per script file, to ACX levels.
1. Speech track: head, pieces, chimes, pauses (digital silence), tail.
2. Room tone: quiet noise (fixed seed) as long as the speech track, mixed underneath everywhere, so there is no digital silence.
3. Gain on the speech track so the finished mix lands on the RMS target; a limiter holds the peak ceiling.
4. MP3, CBR, 44.1 kHz, mono."""
import json
import math
import os
import shutil
import tempfile

import audiomeasure as am
import settings
import voice

FMT = ["-ar", "{sr}", "-ac", "{ch}", "-c:a", "pcm_f32le"]


def fmt(a):
    return [x.format(sr=a["sample_rate"], ch=a["channels"]) for x in FMT]


def make_silence(a, seconds, path):
    am.ffmpeg(["-f", "lavfi", "-i", "anullsrc=r=%d:cl=%s" % (a["sample_rate"], "mono" if a["channels"] == 1 else "stereo"),
               "-t", "%.3f" % seconds] + fmt(a) + [path])


def make_room_tone(a, seconds, path):
    """Very quiet noise (ACX wants room tone, not digital silence). Fixed seed, so output is repeatable."""
    amp = 10 ** (a["room_tone_db"] / 20) * math.sqrt(3)      # uniform white noise: RMS = amplitude / sqrt(3)
    am.ffmpeg(["-f", "lavfi", "-i", "anoisesrc=d=%.3f:c=white:r=%d:a=%.6f:seed=1" % (seconds, a["sample_rate"], amp)] + fmt(a) + [path])


def make_chime(a, level_db, path):
    """A soft sine tone with a smooth rise and fall (sin^2 envelope), peak at level_db."""
    amp = 10 ** (level_db / 20)
    dur = a["chime_s"]
    expr = "%.6f*sin(2*PI*%g*t)*pow(sin(PI*t/%g),2)" % (amp, a["chime_hz"], dur)
    am.ffmpeg(["-f", "lavfi", "-i", "aevalsrc='%s':s=%d:d=%g" % (expr, a["sample_rate"], dur)] + fmt(a) + [path])


def concat(paths, out, tmp):
    lst = os.path.join(tmp, "list_%s.txt" % os.path.basename(out))
    with open(lst, "w") as f:
        for p in paths:
            f.write("file '%s'\n" % p.replace("'", "'\\''"))
    am.ffmpeg(["-f", "concat", "-safe", "0", "-i", lst, "-c", "copy", out])


def build_speech(a, items, chime_db, tmp, tag):
    """Speech track. chime_db None leaves the cue slots silent (used to find the gain before the chime is made)."""
    cache = {}

    def silence(seconds):
        key = "sil%.3f" % seconds
        if key not in cache:
            cache[key] = os.path.join(tmp, key + ".wav")
            make_silence(a, seconds, cache[key])
        return cache[key]

    parts = [silence(a["head_room_tone_s"])]
    for n, it in enumerate(items):
        if "cue" in it:
            if chime_db is None:
                parts.append(silence(a["chime_s"]))
            else:
                if "chime" not in cache:
                    cache["chime"] = os.path.join(tmp, "chime_%s.wav" % tag)
                    make_chime(a, chime_db, cache["chime"])
                parts.append(cache["chime"])
        else:
            p = os.path.join(tmp, "p%05d.wav" % n)
            if not os.path.exists(p):
                am.ffmpeg(["-i", it["wav"]] + fmt(a) + [p])
            parts.append(p)
        if n < len(items) - 1:
            parts.append(silence(it["gap_after_s"]))
    parts.append(silence(a["tail_room_tone_s"]))
    out = os.path.join(tmp, "speech_%s.wav" % tag)
    concat(parts, out, tmp)
    return out


def mix(a, speech, gain, noise, out):
    limit = 10 ** (a["peak_ceiling_db"] / 20)
    am.ffmpeg(["-i", speech, "-i", noise, "-filter_complex",
               "[0:a]volume=%.3fdB,alimiter=limit=%.5f:level=0[s];[s][1:a]amix=inputs=2:normalize=0:duration=first" % (gain, limit)]
              + fmt(a) + [out])


def set_level(a, speech, tmp):
    """Find the gain that puts the finished mix on the RMS target. Returns (mix path, gain dB)."""
    noise = os.path.join(tmp, "noise.wav")
    make_room_tone(a, float(am.probe(speech)["duration"]), noise)
    rms, _ = am.overall(speech)
    gain = a["target_rms_db"] - rms
    for i in range(3):
        out = os.path.join(tmp, "mix_%d.wav" % i)
        mix(a, speech, gain, noise, out)
        got, _ = am.overall(out)
        if abs(got - a["target_rms_db"]) <= 0.3:
            break
        gain += a["target_rms_db"] - got
    return out, gain, noise


def assemble_items(cfg, items, out_mp3):
    """items as written by voice.py. Writes out_mp3. Returns the gain applied to speech in dB."""
    a = dict(cfg["audio"])
    work = settings.resolve(cfg, "work_dir")
    os.makedirs(work, exist_ok=True)
    tmp = tempfile.mkdtemp(prefix="asm_", dir=work)
    try:
        speech = build_speech(a, items, None, tmp, "measure")
        final, gain, noise = set_level(a, speech, tmp)
        if any("cue" in it for it in items):
            # build again with the chime in; it is made so it lands at chime_db after the gain
            speech = build_speech(a, items, a["chime_db"] - gain, tmp, "final")
            final = os.path.join(tmp, "final.wav")
            mix(a, speech, gain, noise, final)
        os.makedirs(os.path.dirname(out_mp3), exist_ok=True)
        am.ffmpeg(["-i", final, "-c:a", "libmp3lame", "-b:a", a["mp3_bitrate"], "-ar", str(a["sample_rate"]), "-ac", str(a["channels"]), out_mp3])
        return gain
    finally:
        shutil.rmtree(tmp, ignore_errors=True)


def run(cfg, book_dir, only=None):
    import bookio
    book = bookio.read_book(book_dir)
    out_dir = settings.audio_dir(cfg, book_dir)
    for name in book["audio_files"]:
        if only and name != only:
            continue
        with open(os.path.join(voice.work_dir_for(cfg, book_dir, name), "manifest.json"), encoding="utf-8") as f:
            items = json.load(f)
        gain = assemble_items(cfg, items, os.path.join(out_dir, name + ".mp3"))
        print("%s.mp3 written (speech gain %.1f dB)" % (name, gain))