ai-music-visualizer/AudioAnalyzer.swift at master · kashviSwami/ai-music-visualizer · GitHub

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
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
import AVFoundation
import Accelerate

class AudioAnalyzer {
    func extractFeatures(from url: URL) -> [String: [Double]]? {
        do {
            let file = try AVAudioFile(forReading: url)
            guard let buffer = AVAudioPCMBuffer(pcmFormat: file.processingFormat, frameCapacity: AVAudioFrameCount(file.length)) else {
                print("Failed to create AVAudioPCMBuffer")
                return nil
            }
            try file.read(into: buffer)
            guard let floatChannelData = buffer.floatChannelData else {
                print("Failed to get float channel data")
                return nil
            }

            let frameLength = Int(buffer.frameLength)
            var data = [Float](repeating: 0.0, count: frameLength)
            for i in 0..<frameLength {
                data[i] = floatChannelData[0][i]
            }

            let mfccs = computeMFCCs(data: data, sampleRate: file.fileFormat.sampleRate)
            let chroma = computeChroma(data: data, sampleRate: file.fileFormat.sampleRate)
            let spectralContrast = computeSpectralContrast(data: data, sampleRate: file.fileFormat.sampleRate)

            return [
                "mfccs": mfccs,
                "chroma": chroma,
                "spectral_contrast": spectralContrast
            ]
        } catch {
            print("Failed to read audio file: \(error.localizedDescription)")
            return nil
        }
    }

    private func computeMFCCs(data: [Float], sampleRate: Double) -> [Double] {
        let fftLength = 1024
        let log2n = vDSP_Length(log2(Double(fftLength)))
        let halfLength = fftLength / 2

        var window = [Float](repeating: 0, count: fftLength)
        var real = [Float](repeating: 0, count: halfLength)
        var imag = [Float](repeating: 0, count: halfLength)
        var splitComplex = DSPSplitComplex(realp: &real, imagp: &imag)

        vDSP_hann_window(&window, vDSP_Length(fftLength), Int32(vDSP_HANN_NORM))
        var windowedData = [Float](repeating: 0, count: fftLength)
        vDSP_vmul(data, 1, window, 1, &windowedData, 1, vDSP_Length(fftLength))

        let fftSetup = vDSP_create_fftsetup(log2n, Int32(FFT_RADIX2))
        windowedData.withUnsafeBufferPointer { dataPointer in
            let dataAddress = dataPointer.baseAddress!
            dataAddress.withMemoryRebound(to: DSPComplex.self, capacity: fftLength) { typeConvertedData in
                vDSP_ctoz(typeConvertedData, 2, &splitComplex, 1, vDSP_Length(halfLength))
            }
        }

        vDSP_fft_zrip(fftSetup!, &splitComplex, 1, log2n, Int32(FFT_FORWARD))
        vDSP_zvabs(&splitComplex, 1, &real, 1, vDSP_Length(halfLength))

        vDSP_destroy_fftsetup(fftSetup)

        let mfccs = real.prefix(13).map { Double($0) }
        return mfccs
    }

    private func computeChroma(data: [Float], sampleRate: Double) -> [Double] {
        let fftLength = 1024
        let log2n = vDSP_Length(log2(Double(fftLength)))
        let halfLength = fftLength / 2

        var window = [Float](repeating: 0, count: fftLength)
        var real = [Float](repeating: 0, count: halfLength)
        var imag = [Float](repeating: 0, count: halfLength)
        var splitComplex = DSPSplitComplex(realp: &real, imagp: &imag)

        vDSP_hann_window(&window, vDSP_Length(fftLength), Int32(vDSP_HANN_NORM))
        var windowedData = [Float](repeating: 0, count: fftLength)
        vDSP_vmul(data, 1, window, 1, &windowedData, 1, vDSP_Length(fftLength))

        let fftSetup = vDSP_create_fftsetup(log2n, Int32(FFT_RADIX2))
        windowedData.withUnsafeBufferPointer { dataPointer in
            let dataAddress = dataPointer.baseAddress!
            dataAddress.withMemoryRebound(to: DSPComplex.self, capacity: fftLength) { typeConvertedData in
                vDSP_ctoz(typeConvertedData, 2, &splitComplex, 1, vDSP_Length(halfLength))
            }
        }

        vDSP_fft_zrip(fftSetup!, &splitComplex, 1, log2n, Int32(FFT_FORWARD))
        vDSP_zvabs(&splitComplex, 1, &real, 1, vDSP_Length(halfLength))

        vDSP_destroy_fftsetup(fftSetup)

        let chroma = real.prefix(12).map { Double($0) }
        return chroma
    }

    private func computeSpectralContrast(data: [Float], sampleRate: Double) -> [Double] {
        let fftLength = 1024
        let log2n = vDSP_Length(log2(Double(fftLength)))
        let halfLength = fftLength / 2

        var window = [Float](repeating: 0, count: fftLength)
        var real = [Float](repeating: 0, count: halfLength)
        var imag = [Float](repeating: 0, count: halfLength)
        var splitComplex = DSPSplitComplex(realp: &real, imagp: &imag)

        vDSP_hann_window(&window, vDSP_Length(fftLength), Int32(vDSP_HANN_NORM))
        var windowedData = [Float](repeating: 0, count: fftLength)
        vDSP_vmul(data, 1, window, 1, &windowedData, 1, vDSP_Length(fftLength))

        let fftSetup = vDSP_create_fftsetup(log2n, Int32(FFT_RADIX2))
        windowedData.withUnsafeBufferPointer { dataPointer in
            let dataAddress = dataPointer.baseAddress!
            dataAddress.withMemoryRebound(to: DSPComplex.self, capacity: fftLength) { typeConvertedData in
                vDSP_ctoz(typeConvertedData, 2, &splitComplex, 1, vDSP_Length(halfLength))
            }
        }

        vDSP_fft_zrip(fftSetup!, &splitComplex, 1, log2n, Int32(FFT_FORWARD))
        vDSP_zvabs(&splitComplex, 1, &real, 1, vDSP_Length(halfLength))

        vDSP_destroy_fftsetup(fftSetup)

        let spectralContrast = real.prefix(6).map { Double($0) }
        return spectralContrast
    }
}