vibesonic/internal/audiofeatures/extract.go

// Package audiofeatures extracts a small set of real, numerically
// computed audio characteristics from a media file — tempo, loudness,
// spectral shape, and a chroma (pitch-class) profile used both for a
// musical key guess and as the basis of a sonic-similarity feature
// vector. Everything here is derived from decoding the actual audio
// with ffmpeg; nothing is inferred from tags, filenames, or genre.
package audiofeatures

import (
	"bytes"
	"context"
	"encoding/binary"
	"fmt"
	"math"
	"os/exec"
	"sort"
)

const (
	sampleRate = 22050
	frameSize  = 2048
	hopSize    = 1024
)

// Features is everything VibeSonic actually managed to measure about a
// track's audio. Fields the analysis couldn't determine with enough
// confidence (e.g. tempo on a clip too short to contain a beat) are
// left at their zero value with the matching Valid flag false, rather
// than filled with a guess.
type Features struct {
	TempoBPM           float64
	TempoValid         bool
	Key                string  // e.g. "C major", "" if undetermined
	LoudnessLUFS       float64 // approximate: mean-square level in dB, not full ITU-R BS.1770
	SpectralCentroidHz float64
	SpectralRolloffHz  float64
	ZeroCrossingRate   float64
	Chroma             [12]float64 // normalized pitch-class energy distribution
}

// Vector returns the fixed-length numeric feature vector used for
// sonic-similarity distance calculations: [tempo, loudness, centroid,
// rolloff, zcr, chroma(12)]. An invalid tempo contributes 0, same as a
// track's other quiet dimensions would if it were literally silent in
// that regard — the similarity index normalizes each dimension across
// the library before comparing, so a shared 0 among untempoed tracks
// doesn't bias results toward or away from any real tempo.
func (f Features) Vector() []float64 {
	tempo := 0.0
	if f.TempoValid {
		tempo = f.TempoBPM
	}
	v := make([]float64, 0, 17)
	v = append(v, tempo, f.LoudnessLUFS, f.SpectralCentroidHz, f.SpectralRolloffHz, f.ZeroCrossingRate)
	v = append(v, f.Chroma[:]...)
	return v
}

// Analyze decodes sourcePath to mono PCM via ffmpeg and computes
// Features from it. It requires at least a few frames of audio
// (roughly 150ms); shorter input still returns loudness/spectral
// values from what's available but leaves tempo unset.
func Analyze(ctx context.Context, ffmpegPath, sourcePath string) (*Features, error) {
	if ffmpegPath == "" {
		ffmpegPath = "ffmpeg"
	}
	pcm, err := decodePCM(ctx, ffmpegPath, sourcePath)
	if err != nil {
		return nil, err
	}
	if len(pcm) == 0 {
		return nil, fmt.Errorf("audiofeatures: no decodable audio in %s", sourcePath)
	}
	return analyzeSamples(pcm), nil
}

func decodePCM(ctx context.Context, ffmpegPath, sourcePath string) ([]float64, error) {
	cmd := exec.CommandContext(ctx, ffmpegPath,
		"-v", "error",
		"-i", sourcePath,
		"-f", "f32le", "-ac", "1", "-ar", fmt.Sprintf("%d", sampleRate),
		"-",
	)
	var stdout, stderr bytes.Buffer
	cmd.Stdout = &stdout
	cmd.Stderr = &stderr
	if err := cmd.Run(); err != nil {
		return nil, fmt.Errorf("ffmpeg decode failed: %w: %s", err, stderr.String())
	}

	raw := stdout.Bytes()
	n := len(raw) / 4
	samples := make([]float64, n)
	for i := 0; i < n; i++ {
		bits := binary.LittleEndian.Uint32(raw[i*4 : i*4+4])
		samples[i] = float64(math.Float32frombits(bits))
	}
	return samples, nil
}

func analyzeSamples(samples []float64) *Features {
	f := &Features{}

	// Overall RMS loudness, in dB relative to full scale — a real
	// (if simplified) measurement of the decoded signal, not a stand-in.
	var sumSq float64
	for _, s := range samples {
		sumSq += s * s
	}
	rms := math.Sqrt(sumSq / float64(len(samples)))
	if rms > 0 {
		f.LoudnessLUFS = 20 * math.Log10(rms)
	} else {
		f.LoudnessLUFS = -120
	}

	var (
		centroidSum, rolloffSum, zcrSum float64
		frameCount                      int
		energyEnvelope                  []float64
		chroma                          [12]float64
	)

	hann := hannWindow(frameSize)
	for start := 0; start+frameSize <= len(samples); start += hopSize {
		frame := make([]float64, frameSize)
		var frameEnergy float64
		var crossings int
		for i := 0; i < frameSize; i++ {
			s := samples[start+i]
			frame[i] = s * hann[i]
			frameEnergy += s * s
			if i > 0 && ((samples[start+i-1] < 0) != (s < 0)) {
				crossings++
			}
		}
		energyEnvelope = append(energyEnvelope, math.Sqrt(frameEnergy/float64(frameSize)))
		zcrSum += float64(crossings) / float64(frameSize)

		mags := magnitudeSpectrum(frame)
		centroid, rolloff := spectralShape(mags, sampleRate, frameSize)
		centroidSum += centroid
		rolloffSum += rolloff
		accumulateChroma(&chroma, mags, sampleRate, frameSize)

		frameCount++
	}

	if frameCount > 0 {
		f.SpectralCentroidHz = centroidSum / float64(frameCount)
		f.SpectralRolloffHz = rolloffSum / float64(frameCount)
		f.ZeroCrossingRate = zcrSum / float64(frameCount)
		normalizeChroma(&chroma)
		f.Chroma = chroma
		f.Key = estimateKey(chroma)
	}

	if bpm, ok := estimateTempo(energyEnvelope, float64(hopSize)/float64(sampleRate)); ok {
		f.TempoBPM = bpm
		f.TempoValid = true
	}

	return f
}

func hannWindow(n int) []float64 {
	w := make([]float64, n)
	for i := range w {
		w[i] = 0.5 * (1 - math.Cos(2*math.Pi*float64(i)/float64(n-1)))
	}
	return w
}

// spectralShape returns the spectral centroid (the "center of mass" of
// the spectrum, in Hz — a brightness measure) and the 85% rolloff
// frequency (below which 85% of the frame's spectral energy lies).
func spectralShape(mags []float64, sampleRate, frameSize int) (centroid, rolloff float64) {
	var weighted, total float64
	binHz := float64(sampleRate) / float64(frameSize)
	for i, m := range mags {
		freq := float64(i) * binHz
		weighted += freq * m
		total += m
	}
	if total == 0 {
		return 0, 0
	}
	centroid = weighted / total

	target := 0.85 * total
	var cum float64
	for i, m := range mags {
		cum += m
		if cum >= target {
			rolloff = float64(i) * binHz
			break
		}
	}
	return centroid, rolloff
}

// accumulateChroma folds a frame's spectral magnitude into 12
// pitch-class bins (assuming standard A440 tuning), skipping the
// sub-audible and very-high-frequency bins where pitch content is
// unreliable.
func accumulateChroma(chroma *[12]float64, mags []float64, sampleRate, frameSize int) {
	binHz := float64(sampleRate) / float64(frameSize)
	for i, m := range mags {
		freq := float64(i) * binHz
		if freq < 27.5 || freq > 5000 || m == 0 {
			continue
		}
		midi := 69 + 12*math.Log2(freq/440)
		pc := int(math.Round(midi)) % 12
		if pc < 0 {
			pc += 12
		}
		chroma[pc] += m
	}
}

func normalizeChroma(chroma *[12]float64) {
	var total float64
	for _, v := range chroma {
		total += v
	}
	if total == 0 {
		return
	}
	for i := range chroma {
		chroma[i] /= total
	}
}

// krumhanslMajor/Minor are the classic Krumhansl-Kessler key profiles:
// the relative perceptual stability of each pitch class within a major
// or minor key, starting from the tonic. Correlating a track's chroma
// vector against all 24 rotations picks the best-fitting key.
var krumhanslMajor = [12]float64{6.35, 2.23, 3.48, 2.33, 4.38, 4.09, 2.52, 5.19, 2.39, 3.66, 2.29, 2.88}
var krumhanslMinor = [12]float64{6.33, 2.68, 3.52, 5.38, 2.60, 3.53, 2.54, 4.75, 3.98, 2.69, 3.34, 3.17}

var pitchClassNames = [12]string{"C", "C♯", "D", "D♯", "E", "F", "F♯", "G", "G♯", "A", "A♯", "B"}

func estimateKey(chroma [12]float64) string {
	type candidate struct {
		name  string
		score float64
	}
	var candidates []candidate
	for tonic := 0; tonic < 12; tonic++ {
		candidates = append(candidates,
			candidate{pitchClassNames[tonic] + " major", correlate(chroma, rotate(krumhanslMajor, tonic))},
			candidate{pitchClassNames[tonic] + " minor", correlate(chroma, rotate(krumhanslMinor, tonic))},
		)
	}
	sort.Slice(candidates, func(i, j int) bool { return candidates[i].score > candidates[j].score })
	if len(candidates) == 0 || candidates[0].score <= 0 {
		return ""
	}
	return candidates[0].name
}

func rotate(profile [12]float64, by int) [12]float64 {
	var out [12]float64
	for i := range out {
		out[i] = profile[(i-by+12)%12]
	}
	return out
}

func correlate(a, b [12]float64) float64 {
	var meanA, meanB float64
	for i := 0; i < 12; i++ {
		meanA += a[i]
		meanB += b[i]
	}
	meanA /= 12
	meanB /= 12
	var num, denomA, denomB float64
	for i := 0; i < 12; i++ {
		da, db := a[i]-meanA, b[i]-meanB
		num += da * db
		denomA += da * da
		denomB += db * db
	}
	if denomA == 0 || denomB == 0 {
		return 0
	}
	return num / math.Sqrt(denomA*denomB)
}

// estimateTempo autocorrelates the frame-wise energy envelope over the
// lag range corresponding to 50-220 BPM and returns the lag with the
// strongest periodicity. It reports ok=false when there isn't enough
// envelope to cover even the slowest candidate tempo, rather than
// returning a meaningless number for a clip too short to have a beat.
func estimateTempo(envelope []float64, hopSeconds float64) (bpm float64, ok bool) {
	minLagSec := 60.0 / 220.0
	maxLagSec := 60.0 / 50.0
	minLag := int(minLagSec / hopSeconds)
	maxLag := int(maxLagSec / hopSeconds)
	if minLag < 1 {
		minLag = 1
	}
	if len(envelope) < maxLag+8 {
		return 0, false
	}

	var mean float64
	for _, v := range envelope {
		mean += v
	}
	mean /= float64(len(envelope))
	centered := make([]float64, len(envelope))
	for i, v := range envelope {
		centered[i] = v - mean
	}

	bestLag := -1
	bestScore := 0.0
	for lag := minLag; lag <= maxLag && lag < len(centered); lag++ {
		var score float64
		for i := 0; i+lag < len(centered); i++ {
			score += centered[i] * centered[i+lag]
		}
		if score > bestScore {
			bestScore = score
			bestLag = lag
		}
	}
	if bestLag <= 0 || bestScore <= 0 {
		return 0, false
	}
	return 60.0 / (float64(bestLag) * hopSeconds), true
}