Extracting audio from video files in Julia

using FFMPEG, WAV, DSP, CairoMakie
set_theme!(theme_dark())
workdir = mkpath(joinpath(tempdir(), "video_audio_demo"))

"/var/folders/y6/plk3_3vd1w35r7vkb_4j6p8h0000gn/T/video_audio_demo"

### Generate a test clip, so the notebook needs no external asset.
### Video: colour bars. Audio: a 440 Hz tone on the left, a 200→4000 Hz chirp on the right —
### distinct per channel, so we can confirm later that the extraction preserved the stereo layout.
clip = joinpath(workdir, "sample.mp4")

ffmpeg_exe(`-y -loglevel error
    -f lavfi -i testsrc2=size=320x240:rate=25:duration=5
    -f lavfi -i "aevalsrc='0.5*sin(2*PI*440*t)|0.5*sin(2*PI*(200*t+380*t*t))':d=5:s=44100:c=stereo"
    -c:v libx264 -pix_fmt yuv420p -c:a aac -b:a 192k -shortest $clip`)

(; path = clip, bytes = filesize(clip))

(path = "/var/folders/y6/plk3_3vd1w35r7vkb_4j6p8h0000gn/T/video_audio_demo/sample.mp4", bytes = 303168)

### Where VideoIO.jl stops: it demuxes the container and *indexes* the audio stream,
### but exports nothing that can decode it.
using VideoIO

av = VideoIO.open(clip)
(; video_streams = av.video_indices,
   audio_streams = av.audio_indices,
   audio_api = filter(n -> occursin("audio", lowercase(String(n))), names(VideoIO)))

### Option 1 — demux to a WAV file, then read it.
### `-vn` drops video; `pcm_s16le` gives uncompressed 16-bit PCM.
wavpath = joinpath(workdir, "sample.wav")
ffmpeg_exe(`-y -loglevel error -i $clip -vn -acodec pcm_s16le $wavpath`)

samples, fs, nbits, _ = wavread(wavpath)
(; file_mb = filesize(wavpath) / 2^20, size = size(samples), fs, nbits, eltype = eltype(samples))

(file_mb = 0.8438243865966797, size = (221184, 2), fs = 44100.0f0, nbits = 0x0010, eltype = Float64)

### Option 2 — skip the intermediate file entirely: ask ffmpeg for raw interleaved PCM
### on stdout and reinterpret the bytes in place. This is the one worth reaching for.
"""
    readaudio(path; fs=44100, channels=2, stream=0)

Decode the audio of any container ffmpeg can read into an `nframes × channels`
`Matrix{Float32}` in [-1, 1]. `stream` selects among multiple audio tracks.
"""
function readaudio(path; fs::Int = 44100, channels::Int = 2, stream::Int = 0)
    bytes = FFMPEG.ffmpeg() do exe
        read(`$exe -loglevel error -i $path -map 0:a:$stream -vn
              -f s16le -acodec pcm_s16le -ac $channels -ar $fs pipe:1`)
    end
    pcm = reinterpret(Int16, bytes)                     # no copy
    frames = permutedims(reshape(pcm, channels, :))     # interleaved → nframes × channels
    Float32.(frames) ./ Float32(typemax(Int16))
end

audio = readaudio(clip)
(; size = size(audio), seconds = size(audio, 1) / 44100, peak = maximum(abs, audio))

(size = (221184, 2), seconds = 5.015510204081632, peak = 0.75838494f0)

### Do the two routes agree? They should be bit-identical — same decoder, same PCM.
wav_f32 = Float32.(samples)
(; same_shape  = size(wav_f32) == size(audio),
   max_absdiff = maximum(abs, wav_f32 .- audio),
   peak_left   = maximum(abs, view(audio, :, 1)),
   peak_right  = maximum(abs, view(audio, :, 2)))

(same_shape = true, max_absdiff = 0.0f0, peak_left = 0.75838494f0, peak_right = 0.51799065f0)

### Did the extraction actually preserve the signal? The left channel should be a flat 440 Hz
### line; the right should ramp 200 → 4000 Hz. Anything else means a wrong sample rate,
### a channel-interleaving mistake, or a byte-order slip.
let fsr = 44100, t = (0:size(audio, 1)-1) ./ fsr, n = 882   # 20 ms
    fig = Figure(size = (900, 560))

    ax1 = Axis(fig[1, 1:2]; title = "waveform (first 20 ms)", xlabel = "time (ms)")
    lines!(ax1, t[1:n] .* 1e3, audio[1:n, 1], label = "L — 440 Hz tone")
    lines!(ax1, t[1:n] .* 1e3, audio[1:n, 2], label = "R — chirp")
    ylims!(ax1, -1.05, 1.05)
    axislegend(ax1; position = :rt, orientation = :horizontal,
               backgroundcolor = (:black, 0.65), framevisible = false)

    local hm
    for (col, name) in enumerate(("left", "right"))
        spec = spectrogram(audio[:, col], 1024, 768; fs = fsr)
        ax = Axis(fig[2, col]; title = "spectrogram — $name",
                  xlabel = "time (s)", ylabel = col == 1 ? "frequency (Hz)" : "")
        hm = heatmap!(ax, DSP.time(spec), DSP.freq(spec),
                      10 .* log10.(power(spec)' .+ eps()))
        ylims!(ax, 0, 5000)
    end
    Colorbar(fig[2, 3], hm; label = "power (dB)")
    fig
end

### Option 3 — no subprocess at all.
###
### "Is there a native package?" — not a pure-Julia one: nothing in the General registry demuxes
### MP4/Matroska (LibSndFile, MP3.jl and Opus.jl all decode *elementary* streams, not containers),
### and a pure-Julia AAC decoder does not exist.
###
### But VideoIO already *links* libavformat/libavcodec and exposes the complete C API as
### `VideoIO.libffmpeg`. So the decode can run in-process via ccall — same library ffmpeg(1)
### uses, no fork/exec, no pipe. What's missing upstream is only the high-level wrapper.
const LF = VideoIO.libffmpeg

function readaudio_inprocess(path)
    fmt = Ref(Ptr{LF.AVFormatContext}(C_NULL))
    LF.avformat_open_input(fmt, path, C_NULL, C_NULL) < 0 && error("could not open $path")
    try
        LF.avformat_find_stream_info(fmt[], C_NULL)
        dec = Ref(Ptr{LF.AVCodec}(C_NULL))
        idx = LF.av_find_best_stream(fmt[], LF.AVMEDIA_TYPE_AUDIO, -1, -1, dec, 0)
        idx < 0 && error("no audio stream")
        par = unsafe_load(unsafe_load(unsafe_load(fmt[].streams), idx + 1).codecpar)
        nch = Int(unsafe_load(par.ch_layout).nb_channels)
        fs  = Int(unsafe_load(par.sample_rate))

        ctx  = LF.avcodec_alloc_context3(dec[])
        cref = Ref(ctx)
        try
            LF.avcodec_parameters_to_context(ctx, par)
            LF.avcodec_open2(ctx, dec[], C_NULL) < 0 && error("could not open decoder")
            # AAC decodes to planar float32, so the frame planes are read directly and
            # libswresample is never needed. Other codecs may hand back a different layout.
            sfmt = LF.AVSampleFormat(unsafe_load(ctx.sample_fmt))
            sfmt == LF.AV_SAMPLE_FMT_FLTP || error("expected planar float32, got $sfmt")

            pkt, frm   = LF.av_packet_alloc(), LF.av_frame_alloc()
            pref, fref = Ref(pkt), Ref(frm)
            chans = [Float32[] for _ in 1:nch]
            try
                while LF.av_read_frame(fmt[], pkt) >= 0
                    if unsafe_load(pkt.stream_index) == idx
                        LF.avcodec_send_packet(ctx, pkt)
                        while LF.avcodec_receive_frame(ctx, frm) == 0
                            n      = Int(unsafe_load(frm.nb_samples))
                            planes = unsafe_load(frm.data)
                            for c in 1:nch
                                append!(chans[c], unsafe_wrap(Array, Ptr{Float32}(planes[c]), n))
                            end
                        end
                    end
                    LF.av_packet_unref(pkt)
                end
            finally
                LF.av_frame_free(fref); LF.av_packet_free(pref)
            end
            return reduce(hcat, chans), fs
        finally
            LF.avcodec_free_context(cref)
        end
    finally
        LF.avformat_close_input(fmt)
    end
end

native, fs_native = readaudio_inprocess(clip)
(; size = size(native), fs = fs_native, peak = maximum(abs, native))

(size = (221184, 2), fs = 44100, peak = 0.7583429f0)

### The in-process route is not merely equivalent — it's slightly *better*.
### `-f s16le` forces the pipe route through a 16-bit requantization (with dither), while
### reading the AVFrame planes keeps the decoder's native float32. The gap is exactly that.
lsb16 = 1 / 32768
(; agree_shape   = size(native) == size(audio),
   max_absdiff   = maximum(abs, native .- audio),
   in_lsb_units  = maximum(abs, native .- audio) / lsb16,
   note          = "sub-2-LSB ⇒ difference is the s16 requantization, not a decode error")

(agree_shape = true, max_absdiff = 5.310774f-5, in_lsb_units = 1.740234375, note = "sub-2-LSB ⇒ difference is the s16 requantization, not a decode error")

### ...and the requantization is avoidable: ask the pipe for `f32le` instead of `s16le` and
### option 2 becomes bit-identical to the in-process decode. So the pipe costs you a subprocess
### and nothing else.
audio_f32 = let bytes = FFMPEG.ffmpeg() do exe
        read(`$exe -loglevel error -i $clip -map 0:a:0 -vn -f f32le -ac 2 -ar 44100 pipe:1`)
    end
    permutedims(reshape(reinterpret(Float32, bytes), 2, :))
end

(; max_absdiff_vs_inprocess = maximum(abs, audio_f32 .- native),
   bit_identical = audio_f32 == native)

(max_absdiff_vs_inprocess = 0.0f0, bit_identical = true)