music-video-gen/flow-state/src/audio/segment.js

// Structural segmentation: where the song changes, and what each part is.
//
// Standard approach — self-similarity plus a checkerboard novelty kernel — kept
// deliberately simple. This is the fuzziest stage in the pipeline, and a wrong
// boundary costs a scene change in a slightly odd place, not a broken video. An
// ML detour here would buy very little.
//
// The one domain-specific trick that earns its keep: boundaries are SNAPPED to
// the bar grid, preferring 4- and 8-bar multiples. Electronic music changes on
// phrase lines, so snapping converts a boundary that is roughly right into one
// that is exactly right.

export const SECTION_KINDS = ['intro', 'build', 'drop', 'sustain', 'breakdown', 'outro'];

const ANALYSIS_HZ = 4;          // coarse grid for the similarity matrix
const KERNEL_SECONDS = 6;       // half-width of the checkerboard kernel
const MIN_SECTION_SECONDS = 12;

function median(values) {
    if (!values.length) return 0;
    const s = Float64Array.from(values).sort();
    return s[Math.floor(s.length / 2)];
}

/** Coarse, L2-normalised feature vectors — the rows of the similarity matrix. */
function buildCoarseVectors(raw, frameCount, fps) {
    const step = Math.max(1, Math.round(fps / ANALYSIS_HZ));
    const count = Math.max(1, Math.floor(frameCount / step));
    const dims = ['bandSub', 'bandLow', 'bandMid', 'bandHigh', 'bandAir', 'centroid', 'flatness'];
    const vectors = [];

    for (let c = 0; c < count; c++) {
        const start = c * step;
        const end = Math.min(frameCount, start + step);
        const v = new Float32Array(dims.length);
        for (let d = 0; d < dims.length; d++) {
            let sum = 0;
            for (let f = start; f < end; f++) sum += raw[dims[d]][f];
            v[d] = sum / Math.max(1, end - start);
        }
        let norm = 0;
        for (let d = 0; d < v.length; d++) norm += v[d] * v[d];
        norm = Math.sqrt(norm) + 1e-9;
        for (let d = 0; d < v.length; d++) v[d] /= norm;
        vectors.push(v);
    }
    return { vectors, step };
}

function cosine(a, b) {
    let dot = 0;
    for (let i = 0; i < a.length; i++) dot += a[i] * b[i];
    return dot;
}

/**
 * Novelty curve. Only the band around the diagonal is needed, so this never
 * materialises the full N×N matrix.
 */
function noveltyCurve(vectors, kernelHalf) {
    const n = vectors.length;
    const novelty = new Float32Array(n);

    // Gaussian-tapered checkerboard weights.
    const size = kernelHalf * 2;
    const weights = new Float32Array(size * size);
    const sigma = kernelHalf / 1.5;
    for (let a = 0; a < size; a++) {
        for (let b = 0; b < size; b++) {
            const da = a - kernelHalf + 0.5;
            const db = b - kernelHalf + 0.5;
            const taper = Math.exp(-(da * da + db * db) / (2 * sigma * sigma));
            const sign = (da * db) > 0 ? 1 : -1;
            weights[a * size + b] = sign * taper;
        }
    }

    for (let i = kernelHalf; i < n - kernelHalf; i++) {
        let sum = 0;
        for (let a = 0; a < size; a++) {
            const ia = i - kernelHalf + a;
            for (let b = 0; b < size; b++) {
                const ib = i - kernelHalf + b;
                sum += weights[a * size + b] * cosine(vectors[ia], vectors[ib]);
            }
        }
        novelty[i] = Math.max(0, sum);
    }

    let peak = 0;
    for (let i = 0; i < n; i++) if (novelty[i] > peak) peak = novelty[i];
    if (peak > 1e-9) for (let i = 0; i < n; i++) novelty[i] /= peak;
    return novelty;
}

function pickPeaks(novelty, minDistance, threshold) {
    const peaks = [];
    const n = novelty.length;
    for (let i = 1; i < n - 1; i++) {
        if (novelty[i] < threshold) continue;
        if (novelty[i] < novelty[i - 1] || novelty[i] < novelty[i + 1]) continue;
        // Local maximum within the exclusion window.
        let isMax = true;
        const lo = Math.max(0, i - minDistance);
        const hi = Math.min(n - 1, i + minDistance);
        for (let j = lo; j <= hi; j++) if (novelty[j] > novelty[i]) { isMax = false; break; }
        if (!isMax) continue;
        if (peaks.length && i - peaks[peaks.length - 1] < minDistance) {
            if (novelty[i] > novelty[peaks[peaks.length - 1]]) peaks[peaks.length - 1] = i;
            continue;
        }
        peaks.push(i);
    }
    return peaks;
}

/**
 * Snap a time to the musical grid. Prefers, in order: an 8-bar line, a 4-bar
 * line, then any downbeat — but only if one is close enough to be plausibly the
 * same boundary.
 */
function snapToGrid(time, downbeats, barsPerPhrase = 8) {
    if (!downbeats.length) return time;
    const tolerance = 2.0;

    const tryLevels = [barsPerPhrase, 4, 1];
    for (const level of tryLevels) {
        let best = null;
        let bestDist = Infinity;
        for (let i = 0; i < downbeats.length; i += level) {
            const d = Math.abs(downbeats[i] - time);
            if (d < bestDist) { bestDist = d; best = downbeats[i]; }
        }
        if (best !== null && bestDist <= tolerance) return best;
    }
    return time;
}

function linearSlope(values) {
    const n = values.length;
    if (n < 2) return 0;
    let sx = 0, sy = 0, sxy = 0, sxx = 0;
    for (let i = 0; i < n; i++) {
        sx += i; sy += values[i]; sxy += i * values[i]; sxx += i * i;
    }
    const denom = n * sxx - sx * sx;
    if (Math.abs(denom) < 1e-12) return 0;
    return ((n * sxy - sx * sy) / denom) * n; // normalised to "change across the section"
}

/**
 * Label a section from its own statistics plus where it sits in the track.
 * Thresholds are relative to the track, never absolute, so a quiet ambient piece
 * still gets a full set of labels rather than being classified as one long intro.
 */
function classify(section, context) {
    const { energy, slope, flux, index, count, endEnergy } = section;
    const { medianEnergy, isFirst, isLast } = context;

    const low = energy < medianEnergy * 0.72;
    const high = energy > medianEnergy * 1.12;

    if (isFirst && energy < medianEnergy) return 'intro';
    if (isLast && (low || slope < -0.15)) return 'outro';
    if (slope > 0.18 && endEnergy > medianEnergy * 0.95) return 'build';
    if (high && flux > context.medianFlux * 0.95) return 'drop';
    if (low) return 'breakdown';
    if (index === 0) return 'intro';
    if (index === count - 1) return 'outro';
    return 'sustain';
}

/**
 * @returns {Array<{index,start,end,startFrame,endFrame,kind,energy,slope,flux,centroid}>}
 */
export function segment(raw, frameCount, fps, tempo) {
    const { vectors, step } = buildCoarseVectors(raw, frameCount, fps);
    const kernelHalf = Math.max(4, Math.round(KERNEL_SECONDS * ANALYSIS_HZ));

    let boundaryTimes = [];
    if (vectors.length > kernelHalf * 2 + 4) {
        const novelty = noveltyCurve(vectors, kernelHalf);
        const minDistance = Math.round(MIN_SECTION_SECONDS * ANALYSIS_HZ);

        // Adaptive threshold: mean + a fraction of the spread, so a track with
        // gentle transitions still yields boundaries and a busy one isn't shredded.
        let mean = 0;
        for (let i = 0; i < novelty.length; i++) mean += novelty[i];
        mean /= novelty.length;
        let variance = 0;
        for (let i = 0; i < novelty.length; i++) variance += (novelty[i] - mean) ** 2;
        const std = Math.sqrt(variance / novelty.length);
        const threshold = mean + std * 0.6;

        boundaryTimes = pickPeaks(novelty, minDistance, threshold)
            .map((i) => (i * step) / fps);
    }

    // Snap to the bar grid, dedupe, and drop anything too close to the edges.
    const duration = frameCount / fps;
    const snapped = boundaryTimes
        .map((t) => snapToGrid(t, tempo.downbeats))
        .filter((t) => t > MIN_SECTION_SECONDS * 0.5 && t < duration - MIN_SECTION_SECONDS * 0.5)
        .sort((a, b) => a - b);

    const bounds = [0];
    for (const t of snapped) {
        if (t - bounds[bounds.length - 1] >= MIN_SECTION_SECONDS) bounds.push(t);
    }
    bounds.push(duration);

    // Build sections and gather their statistics.
    const sections = [];
    for (let i = 0; i < bounds.length - 1; i++) {
        const start = bounds[i];
        const end = bounds[i + 1];
        const startFrame = Math.round(start * fps);
        const endFrame = Math.min(frameCount, Math.round(end * fps));
        const n = Math.max(1, endFrame - startFrame);

        let energySum = 0, fluxSum = 0, centroidSum = 0, flatnessSum = 0, widthSum = 0;
        const coarseEnergy = [];
        const coarseStep = Math.max(1, Math.floor(n / 24));
        for (let f = startFrame; f < endFrame; f++) {
            energySum += raw.loudness[f];
            fluxSum += raw.flux[f];
            centroidSum += raw.centroid[f];
            flatnessSum += raw.flatness[f];
            widthSum += raw.width[f];
            if ((f - startFrame) % coarseStep === 0) coarseEnergy.push(raw.loudness[f]);
        }

        const tailStart = Math.max(startFrame, endFrame - Math.round(fps * 4));
        let endEnergy = 0;
        for (let f = tailStart; f < endFrame; f++) endEnergy += raw.loudness[f];
        endEnergy /= Math.max(1, endFrame - tailStart);

        sections.push({
            index: i,
            start, end,
            startFrame, endFrame,
            duration: end - start,
            energy: energySum / n,
            flux: fluxSum / n,
            centroid: centroidSum / n,
            flatness: flatnessSum / n,
            width: widthSum / n,
            slope: linearSlope(coarseEnergy),
            endEnergy,
            kind: 'sustain',
        });
    }

    const context = {
        medianEnergy: median(sections.map((s) => s.energy)) || 1e-6,
        medianFlux: median(sections.map((s) => s.flux)) || 1e-6,
    };
    sections.forEach((s, i) => {
        s.kind = classify(
            { ...s, index: i, count: sections.length },
            { ...context, isFirst: i === 0, isLast: i === sections.length - 1 },
        );
    });

    return sections;
}