-
Notifications
You must be signed in to change notification settings - Fork 0
/
spectrogram.go
95 lines (85 loc) · 2.56 KB
/
spectrogram.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
package spectrogram
import (
"errors"
"github.com/mjibson/go-dsp/fft"
"github.com/mjibson/go-dsp/window"
"github.com/ngyewch/go-spectrogram/pkg/audio"
"math"
)
type WindowFunction func(int) []float64
type SpectrogramOptions struct {
Channel uint
FftSamples uint
Overlap *uint
Segments *uint
WindowFunction WindowFunction
}
type Spectrogram struct {
SampleRate uint
NumChannels uint
FftSamples uint
Data [][]float64
}
func GenerateSpectrogram(audioFile audio.Source, options SpectrogramOptions) (*Spectrogram, error) {
info := audioFile.Info()
frames := audioFile.Frames()
channel := int(options.Channel)
if channel < 0 || channel >= info.NumChannels {
return nil, errors.New("invalid channel number")
}
fftSamples := options.FftSamples
if !IsPowerOfTwo(fftSamples) {
return nil, errors.New("fftSamples must be a power of 2")
}
overlap := 0
hop := 0
if (options.Segments != nil) && (options.Overlap != nil) {
return nil, errors.New("cannot specify both Segments and Overlap")
} else if (options.Segments == nil) && (options.Overlap == nil) {
overlap = 0
hop = int(fftSamples)
} else if (options.Segments == nil) && (options.Overlap != nil) {
if *options.Overlap >= fftSamples {
return nil, errors.New("overlap must be less than fftSamples")
}
overlap = int(*options.Overlap)
hop = int(fftSamples) - overlap
} else if (options.Segments != nil) && (options.Overlap == nil) {
if *options.Segments <= 1 {
return nil, errors.New("segments must be greater than 1")
}
hop = (len(frames) - int(fftSamples)) / int(*options.Segments-1)
if hop < 1 {
hop = 1
} else if hop > int(fftSamples) {
hop = int(fftSamples)
}
overlap = int(fftSamples) - hop
}
bSi := 2 / float64(fftSamples)
buffer := make([]float64, int(fftSamples))
specColumns := make([][]float64, 0)
for i := 0; (i + int(fftSamples)) <= len(frames); i += hop {
start := i
end := start + int(fftSamples)
currentFrame := frames[start:end]
for j := 0; j < int(fftSamples); j++ {
buffer[j] = currentFrame[j][channel]
}
window.Apply(buffer, options.WindowFunction)
fftResult := fft.FFTReal(buffer)
specColumn := make([]float64, len(buffer)/2)
for j := 0; j < len(buffer)/2; j++ {
val := fftResult[j]
mag := math.Sqrt((real(val)*real(val))+(imag(val)*imag(val))) * bSi
specColumn[j] = 20 * math.Log10(mag)
}
specColumns = append(specColumns, specColumn)
}
return &Spectrogram{
SampleRate: uint(info.SampleRate),
NumChannels: uint(info.NumChannels),
FftSamples: fftSamples,
Data: specColumns,
}, nil
}