scene-actor-extraction/src/inference/audio_logpsd.hpp

#pragma once
// Per-second audio log-PSD, C++ parity with scripts/scene_detector/
// extract_audio_features.py — the audio tower input for the XGBoost scene
// detector. Decodes the whole track to mono 16 kHz, then one FFT per second over
// a 4 s Hann-windowed window, power pooled into geomspace log-frequency bands,
// L1-normalised (shape not loudness) and log1p-compressed.
//
// Must match the Python exactly (SR=16000, WIN_SEC=4, N_BINS=64→geomspace unique
// edges, log1p(band*1e3)); the shipped model was trained on those features.

extern "C" {
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
#include <libavutil/opt.h>
#include <libswresample/swresample.h>
}
#include <fftw3.h>

#include <algorithm>
#include <cmath>
#include <cstdint>
#include <stdexcept>
#include <string>
#include <vector>

class AudioLogPSD {
public:
    static constexpr int   kSR = 16000;
    static constexpr double kHop = 1.0;    // 1 feature row / second
    static constexpr double kWin = 4.0;    // FFT window seconds
    static constexpr int   kNBins = 64;    // geomspace target (dedups to ~57)

    // Returns [T][B] per-second log-PSD (T ≈ film seconds, B ≈ 57), aligned to the
    // 1 fps grid. Empty on decode failure (caller then feeds a zero block).
    static std::vector<std::vector<float>> extract(const std::string& path) {
        std::vector<float> mono = decode_mono_16k(path);
        if (mono.empty()) return {};
        return features(mono);
    }

    // Public for the parity harness.
    static std::vector<std::vector<float>> features(const std::vector<float>& mono) {
        const int win = int(kSR * kWin), hop = int(kSR * kHop);
        const int T = int(mono.size()) / hop;
        if (T <= 0) return {};
        const int nfreq = win/2 + 1;
        std::vector<int> edges = geomspace_edges(nfreq);
        const int nb = int(edges.size()) - 1;

        // Hann window (matches scipy.signal.windows.hann, sym=True default → but
        // numpy code uses sps.windows.hann(win) which is symmetric).
        std::vector<double> hann(win);
        for (int i = 0; i < win; ++i)
            hann[i] = 0.5 - 0.5*std::cos(2.0*M_PI*i/(win-1));

        std::vector<double> in(win);
        auto* out = fftw_alloc_complex(nfreq);
        fftw_plan plan = fftw_plan_dft_r2c_1d(win, in.data(), out, FFTW_ESTIMATE);

        std::vector<std::vector<float>> feat(T, std::vector<float>(nb, 0.f));
        const int half = win/2;
        for (int t = 0; t < T; ++t) {
            int centre = t*hop + hop/2;
            int s = centre - half;
            for (int i = 0; i < win; ++i) {
                int idx = s + i;
                double v = (idx >= 0 && idx < int(mono.size())) ? mono[idx] : 0.0;
                in[i] = v * hann[i];
            }
            fftw_execute(plan);
            // power spectrum + 1e-12
            std::vector<double> psd(nfreq);
            for (int i = 0; i < nfreq; ++i)
                psd[i] = out[i][0]*out[i][0] + out[i][1]*out[i][1] + 1e-12;
            std::vector<double> band(nb, 0.0);
            double tot = 0.0;
            for (int b = 0; b < nb; ++b) {
                for (int i = edges[b]; i < edges[b+1]; ++i) band[b] += psd[i];
                tot += band[b];
            }
            for (int b = 0; b < nb; ++b)
                feat[t][b] = float(std::log1p(band[b]/tot * 1e3));
        }
        fftw_destroy_plan(plan); fftw_free(out);
        return feat;
    }

private:
    // np.unique(np.geomspace(1, nfreq-1, N_BINS+1).astype(int))
    static std::vector<int> geomspace_edges(int nfreq) {
        const int n = kNBins + 1;
        double a = std::log(1.0), b = std::log(double(nfreq-1));
        std::vector<int> raw(n);
        for (int i = 0; i < n; ++i)
            raw[i] = int(std::exp(a + (b-a)*i/(n-1)));   // .astype(int) truncates
        std::vector<int> uniq;
        for (int v : raw) if (uniq.empty() || v != uniq.back()) uniq.push_back(v);
        return uniq;
    }

    static std::vector<float> decode_mono_16k(const std::string& path) {
        AVFormatContext* fmt = nullptr;
        if (avformat_open_input(&fmt, path.c_str(), nullptr, nullptr) < 0) return {};
        std::vector<float> out;
        SwrContext* swr = nullptr; AVCodecContext* dec = nullptr;
        AVPacket* pkt = av_packet_alloc(); AVFrame* fr = av_frame_alloc();
        try {
            if (avformat_find_stream_info(fmt, nullptr) < 0) throw 0;
            int ai = av_find_best_stream(fmt, AVMEDIA_TYPE_AUDIO, -1, -1, nullptr, 0);
            if (ai < 0) throw 0;
            AVStream* st = fmt->streams[ai];
            const AVCodec* codec = avcodec_find_decoder(st->codecpar->codec_id);
            dec = avcodec_alloc_context3(codec);
            avcodec_parameters_to_context(dec, st->codecpar);
            if (avcodec_open2(dec, codec, nullptr) < 0) throw 0;

            AVChannelLayout out_ch = AV_CHANNEL_LAYOUT_MONO;
            swr_alloc_set_opts2(&swr, &out_ch, AV_SAMPLE_FMT_FLT, kSR,
                                &dec->ch_layout, dec->sample_fmt,
                                dec->sample_rate ? dec->sample_rate : kSR, 0, nullptr);
            if (!swr || swr_init(swr) < 0) throw 0;

            while (av_read_frame(fmt, pkt) >= 0) {
                if (pkt->stream_index == ai && avcodec_send_packet(dec, pkt) >= 0) {
                    while (avcodec_receive_frame(dec, fr) >= 0) {
                        int max_out = swr_get_out_samples(swr, fr->nb_samples);
                        size_t base = out.size(); out.resize(base + max_out);
                        uint8_t* dst = reinterpret_cast<uint8_t*>(out.data() + base);
                        int got = swr_convert(swr, &dst, max_out,
                            (const uint8_t**)fr->extended_data, fr->nb_samples);
                        out.resize(base + std::max(0, got));
                    }
                }
                av_packet_unref(pkt);
            }
        } catch (...) { out.clear(); }
        if (swr) swr_free(&swr);
        if (dec) avcodec_free_context(&dec);
        av_frame_free(&fr); av_packet_free(&pkt);
        avformat_close_input(&fmt);
        return out;
    }
};