songsee/internal/dsp/spectrogram.go
2026-01-02 14:38:42 +01:00

101 lines
2.0 KiB
Go

// Package dsp provides spectral analysis utilities.
package dsp
import (
"math"
)
// Spectrogram contains log-magnitude FFT frames.
type Spectrogram struct {
Frames int
Bins int
Values []float64
Min float64
Max float64
SampleRate int
WindowSize int
HopSize int
BinHz float64
}
// HannWindow returns a Hann window of length n.
func HannWindow(n int) []float64 {
w := make([]float64, n)
if n == 1 {
w[0] = 1
return w
}
for i := 0; i < n; i++ {
w[i] = 0.5 - 0.5*math.Cos(2*math.Pi*float64(i)/float64(n-1))
}
return w
}
// ComputeSpectrogram computes a log-magnitude spectrogram.
func ComputeSpectrogram(samples []float64, sampleRate, windowSize, hopSize int) Spectrogram {
if windowSize <= 0 {
windowSize = 2048
}
if hopSize <= 0 {
hopSize = windowSize / 4
}
if hopSize <= 0 {
hopSize = 1
}
if sampleRate <= 0 {
sampleRate = 44100
}
frames := 1
if len(samples) > windowSize {
frames = 1 + (len(samples)-windowSize+hopSize-1)/hopSize
}
bins := windowSize/2 + 1
values := make([]float64, frames*bins)
window := HannWindow(windowSize)
minVal := math.Inf(1)
maxVal := math.Inf(-1)
eps := 1e-9
frame := make([]complex128, windowSize)
for f := 0; f < frames; f++ {
start := f * hopSize
for i := 0; i < windowSize; i++ {
idx := start + i
if idx < len(samples) {
frame[i] = complex(samples[idx]*window[i], 0)
} else {
frame[i] = 0
}
}
FFTInPlace(frame)
for b := 0; b < bins; b++ {
re := real(frame[b])
im := imag(frame[b])
mag := math.Sqrt(re*re + im*im)
db := 20 * math.Log10(mag+eps)
values[f*bins+b] = db
if db < minVal {
minVal = db
}
if db > maxVal {
maxVal = db
}
}
}
binHz := float64(sampleRate) / float64(windowSize)
return Spectrogram{
Frames: frames,
Bins: bins,
Values: values,
Min: minVal,
Max: maxVal,
SampleRate: sampleRate,
WindowSize: windowSize,
HopSize: hopSize,
BinHz: binHz,
}
}