yum-archive/TaSTT-Whisper

High-performance GPGPU inference of OpenAI's Whisper automatic speech recognition (ASR) model

git clone https://git.yummers.dev/yum-archive/TaSTT-Whisper

yumFix audio normalization5929750

master
branchesmasterall refs →
2.9 KiB116 linesraw
1#pragma once
2#include <algorithm>
3#include <fstream>
4#include <vector>
5
6namespace Whisper
7{
8	struct AudioBuffer
9	{
10		std::vector<float> mono;
11		std::vector<float> stereo;
12
13		void appendMono( const float* rsi, size_t countFloats );
14		void appendDownmixedStereo( const float* rsi, size_t countFloats );
15		void appendStereo( const float* rsi, size_t countFloats );
16
17		using pfnAppendSamples = void( AudioBuffer::* )( const float* rsi, size_t countFloats );
18
19		inline static pfnAppendSamples appendSamplesFunc( bool sourceMono, bool wantStereo )
20		{
21			if( sourceMono )
22				return &AudioBuffer::appendMono;
23			else if( !wantStereo )
24				return &AudioBuffer::appendDownmixedStereo;
25			else
26				return &AudioBuffer::appendStereo;
27		}
28
29		void clear()
30		{
31			mono.clear();
32			stereo.clear();
33		}
34
35		void swap( AudioBuffer& that )
36		{
37			mono.swap( that.mono );
38			stereo.swap( that.stereo );
39		}
40
41		void resize( size_t len )
42		{
43			assert( len <= mono.size() );
44			mono.resize( len );
45			if( !stereo.empty() )
46				stereo.resize( len * 2 );
47		}
48
49		void dropFirst(size_t len)
50		{
51			if (len >= mono.size()) {
52				mono.clear();
53				return;
54			}
55			size_t remainder = mono.size() - len;
56			auto tmp = std::vector<float>(remainder);
57			memcpy(tmp.data(), mono.data() + len, remainder);
58			mono = std::move(tmp);
59		}
60
61		void retainLast(size_t len)
62		{
63			if (len >= mono.size()) {
64				return;
65			}
66			size_t prefix_len = mono.size() - len;
67			auto tmp = std::vector<float>(len);
68			memcpy(tmp.data(), mono.data() + prefix_len, len);
69			mono = std::move(tmp);
70		}
71
72		void normalize()
73		{
74			const auto &min = *std::min_element(mono.begin(), mono.end());
75			const auto &max = *std::max_element(mono.begin(), mono.end());
76
77			for (auto& elm : mono) {
78				elm -= min;
79				elm /= (max - min) + 1;
80			}
81		}
82
83		void save(const char* path, const int sample_rate) {
84			const int n_samples = mono.size();
85			const int bits_per_sample = sizeof(mono[0]) * 8;
86			const int n_channels = 1;
87			const int byte_rate = sample_rate * n_channels * bits_per_sample / 8;
88			const int block_align = n_channels * bits_per_sample / 8;
89			const int data_chunk_size = n_samples * n_channels * bits_per_sample / 8;
90			const int file_size = 36 + data_chunk_size;
91
92			std::ofstream ofs(path, std::ios::out | std::ios::binary);
93			ofs.write("RIFF", 4);
94			ofs.write((char*)&file_size, 4);
95			ofs.write("WAVE", 4);
96
97			ofs.write("fmt ", 4);
98			const int fmt_chunk_size = 16;
99			ofs.write((char*)&fmt_chunk_size, 4);
100			const short audio_format = 1;  // PCM
101			ofs.write((char*)&audio_format, 2);
102			ofs.write((char*)&n_channels, 2);
103			ofs.write((char*)&sample_rate, 4);
104			ofs.write((char*)&byte_rate, 4);
105			ofs.write((char*)&block_align, 2);
106			ofs.write((char*)&bits_per_sample, 2);
107
108			ofs.write("data", 4);
109			ofs.write((char*)&data_chunk_size, 4);
110			for (int i = 0; i < n_samples; i++) {
111				short sample = (short)(mono[i] * 32767.0f);
112				ofs.write((char*)&sample, 2);
113			}
114		};
115	};
116}