Make an SDL/OpenGL-frontend; might be more intuitive to read, although I'm not sure...
[pitch] / pitchdetector.cpp
1 #include <stdio.h>
2 #include <string.h>
3 #include <stdlib.h>
4 #include <unistd.h>
5 #include <fcntl.h>
6 #include <complex>
7 #include <cassert>
8 #include <algorithm>
9 #include <fftw3.h>
10 #include "pitchdetector.h"
11
12 PitchDetector::PitchDetector(unsigned sample_rate, unsigned fft_length, unsigned pad_factor, unsigned overlap)
13         : sample_rate(sample_rate), fft_length(fft_length), pad_factor(pad_factor), overlap(overlap)
14 {
15         in = reinterpret_cast<double *> (fftw_malloc(sizeof(double) * fft_length / pad_factor));
16         in_windowed = reinterpret_cast<double *> (fftw_malloc(sizeof(double) * fft_length));
17         out = reinterpret_cast<std::complex<double> *> (fftw_malloc(sizeof(std::complex<double>) * (fft_length / 2 + 1)));
18         bins = reinterpret_cast<double *> (fftw_malloc(sizeof(double) * (fft_length / 2 + 1)));
19
20         memset(in, 0, sizeof(double) * fft_length / pad_factor);
21
22         plan = fftw_plan_dft_r2c_1d(fft_length, in_windowed, reinterpret_cast<fftw_complex *> (out), FFTW_ESTIMATE);
23         
24         // Initialize the Hamming window
25         window_data = new double[fft_length / pad_factor];
26         for (unsigned i = 0; i < fft_length / pad_factor; ++i) {
27                 window_data[i] = 0.54 - 0.46 * cos(2.0 * M_PI * double(i) / double(fft_length/pad_factor - 1));
28         }
29 }
30
31 PitchDetector::~PitchDetector()
32 {
33         fftw_free(in);
34         fftw_free(in_windowed);
35         fftw_free(out);
36         fftw_free(bins);
37 }
38
39 std::pair<double, double> PitchDetector::detect_pitch(short *buf)
40 {
41         unsigned buf_len = fft_length / pad_factor / overlap;
42         memmove(in, in + buf_len, (fft_length / pad_factor - buf_len) * sizeof(double));
43         
44         for (unsigned i = 0; i < buf_len; ++i)
45                 in[i + (fft_length / pad_factor - buf_len)] = double(buf[i]);
46
47         apply_window(in, in_windowed, fft_length);
48         fftw_execute(plan);
49         find_peak_magnitudes(out, bins, fft_length);
50         std::pair<double, double> peak = find_peak(bins, fft_length);
51         if (peak.first > 0.0)
52                 peak = adjust_for_overtones(peak, bins, fft_length);
53
54         return peak;
55 }
56
57 // Apply a standard Hamming window to our input data.
58 void PitchDetector::apply_window(double *in, double *out, unsigned num_samples)
59 {
60         for (unsigned i = 0; i < num_samples / pad_factor; ++i) {
61                 out[i] = in[i] * window_data[i];
62         }
63         for (unsigned i = num_samples / pad_factor; i < num_samples; ++i) {
64                 out[i] = 0.0;
65         }
66 }
67
68 void PitchDetector::find_peak_magnitudes(std::complex<double> *in, double *out, unsigned num_samples)
69 {
70         for (unsigned i = 0; i < num_samples / 2 + 1; ++i)
71                 out[i] = abs(in[i]);
72 }
73
74 std::pair<double, double> PitchDetector::find_peak(double *in, unsigned num_samples)
75 {
76         double best_peak = in[5];
77         unsigned best_bin = 5;
78
79         for (unsigned i = 6; i < num_samples / 2 + 1; ++i) {
80                 if (in[i] > best_peak) {
81                         best_peak = in[i];
82                         best_bin = i;
83                 }
84 #if 0
85                 if (20.0 * log10(in[i] / fft_length) > 0.0) {
86                         printf("PEAK: %+4.2f dB   %5.2f Hz\n",
87                                 20.0 * log10(in[i] / fft_length),
88                                 bin_to_freq(i, num_samples));
89                 }
90 #endif
91         }
92         
93         if (best_bin == 0 || best_bin == num_samples / 2) {
94                 return std::make_pair(-1.0, 0.0);
95         }
96
97 #if 0
98         printf("undertone strength: %+4.2f %+4.2f %+4.2f [%+4.2f] %+4.2f %+4.2f %+4.2f\n",
99                 20.0 * log10(in[best_bin*4] / fft_length),
100                 20.0 * log10(in[best_bin*3] / fft_length),
101                 20.0 * log10(in[best_bin*2] / fft_length),
102                 20.0 * log10(in[best_bin] / fft_length),
103                 20.0 * log10(in[best_bin/2] / fft_length),
104                 20.0 * log10(in[best_bin/3] / fft_length),
105                 20.0 * log10(in[best_bin/4] / fft_length));
106 #endif
107
108         // see if we might have hit an overtone (set a limit of 10dB)
109         for (unsigned i = 6; i >= 1; --i) {
110                 if (best_bin != best_bin / i &&
111                     20.0 * log10(in[best_bin] / in[best_bin / i]) < 10.0f &&
112                     best_bin / i >= 5) {
113 #if 0
114                         printf("Overtone of degree %u!\n", i);
115 #endif
116                         best_bin /= i;
117
118                         // consider sliding one bin up or down
119                         if (best_bin > 1 && in[best_bin - 1] > in[best_bin] && in[best_bin - 1] > in[best_bin - 2]) {
120                                 --best_bin;
121                         } else if (best_bin < num_samples / 2 - 1 && in[best_bin + 1] > in[best_bin] && in[best_bin + 1] > in[best_bin + 2]) {
122                                 ++best_bin;
123                         }
124                            
125                         break;
126                 }
127         }
128
129         if (best_bin == 0 || best_bin == num_samples / 2) {
130                 return std::make_pair(-1.0, 0.0);
131         }
132         std::pair<double, double> peak = 
133                 interpolate_peak(in[best_bin - 1],
134                                  in[best_bin],
135                                  in[best_bin + 1]);
136                 
137         return std::make_pair(bin_to_freq(double(best_bin) + peak.first, num_samples), peak.second);
138 }
139
140 // it's perhaps not ideal to _first_ find the peak and _then_ the harmonics --
141 // ideally, one should find the set of all peaks and then determine the likely
142 // base from that... something for later. :-)
143 std::pair<double, double> PitchDetector::adjust_for_overtones(std::pair<double, double> base, double *in, unsigned num_samples)
144 {
145         double mu = base.first, var = 1.0 / (base.second * base.second);
146                 
147         //printf("Base at %.2f (amp=%5.2fdB)\n", base.first, base.second);
148
149         for (unsigned i = 2; i < 10; ++i) {
150                 unsigned middle = unsigned(floor(freq_to_bin(base.first, num_samples) * i + 0.5));
151                 unsigned lower = middle - (i+1)/2, upper = middle + (i+1)/2;
152
153                 if (lower < 1)
154                         lower = 1;
155                 if (upper >= num_samples)
156                         upper = num_samples - 2;
157
158                 // printf("Searching in [%u,%u] = %f..%f\n", lower, upper, bin_to_freq(lower, num_samples), bin_to_freq(upper, num_samples));
159
160                 // search for a peak in this interval
161                 double best_harmonics_freq = -1.0;
162                 double best_harmonics_amp = -1.0;
163                 for (unsigned j = lower; j <= upper; ++j) {
164                         if (in[j] > in[j-1] && in[j] > in[j+1]) {
165                                 std::pair<double, double> peak =
166                                         interpolate_peak(in[j - 1],
167                                                 in[j],
168                                                 in[j + 1]);
169                                 
170                                 if (peak.second > best_harmonics_amp) {
171                                         best_harmonics_freq = bin_to_freq(j + peak.first, num_samples);
172                                         best_harmonics_amp = peak.second;
173                                 }
174                         }
175                 }
176
177                 if (best_harmonics_amp <= 0.0)
178                         continue;
179
180                 //printf("Found overtone %u at %.2f (amp=%5.2fdB)\n", i, best_harmonics_freq,
181                 //       best_harmonics_amp);
182
183                 double this_mu = best_harmonics_freq / double(i);
184                 double this_var = 1.0 / (best_harmonics_amp * best_harmonics_amp);
185
186                 double k = var / (var + this_var);
187                 mu = (1.0 - k) * mu + k * this_mu;
188                 var *= (1.0 - k);
189         }
190         return std::make_pair(mu, base.second);
191 }
192
193 double PitchDetector::bin_to_freq(double bin, unsigned num_samples)
194 {
195         return bin * sample_rate / double(num_samples);
196 }
197 double PitchDetector::freq_to_bin(double freq, unsigned num_samples)
198 {
199         return freq * double(num_samples) / double(sample_rate);
200 }
201
202 /*
203  * Given three bins, find the interpolated real peak based
204  * on their magnitudes. To do this, we execute the following
205  * plan:
206  * 
207  * Fit a polynomial of the form ax^2 + bx + c = 0 to the data
208  * we have. Maple readily yields our coefficients, assuming
209  * that we have the values at x=-1, x=0 and x=1:
210  *
211  * > f := x -> a*x^2 + b*x + c;                                
212  * 
213  *                               2
214  *           f := x -> a x  + b x + c
215  * 
216  * > cf := solve({ f(-1) = ym1, f(0) = y0, f(1) = y1 }, { a, b, c });
217  * 
218  *                            y1    ym1       y1    ym1
219  *        cf := {c = y0, b = ---- - ---, a = ---- + --- - y0}
220  *                            2      2        2      2
221  *
222  * Now let's find the maximum point for the polynomial (it has to be
223  * a maximum, since y0 is the greatest value):
224  *
225  * > xmax := solve(subs(cf, diff(f(x), x)) = 0, x);
226  * 
227  *                                -y1 + ym1
228  *                   xmax := -------------------
229  *                           2 (y1 + ym1 - 2 y0)
230  *
231  * We could go further and insert {fmax,a,b,c} into the original
232  * polynomial, but it just gets hairy. We calculate a, b and c separately
233  * instead.
234  *
235  * http://www-ccrma.stanford.edu/~jos/parshl/Peak_Detection_Steps_3.html
236  * claims that detection is almost twice as good when using dB scale instead
237  * of linear scale for the input values, so we use that. (As a tiny bonus,
238  * we get back dB scale from the function.)
239  */
240 std::pair<double, double> PitchDetector::interpolate_peak(double ym1, double y0, double y1)
241 {
242         ym1 = log10(ym1);
243         y0 = log10(y0);
244         y1 = log10(y1);
245
246 #if 0
247         assert(y0 >= y1);
248         assert(y0 >= ym1);
249 #endif
250         
251         double a = 0.5 * y1 + 0.5 * ym1 - y0;
252         double b = 0.5 * y1 - 0.5 * ym1;
253         double c = y0;
254         
255         double xmax = (ym1 - y1) / (2.0 * (y1 + ym1 - 2.0 * y0));
256         double ymax = 20.0 * (a * xmax * xmax + b * xmax + c) - 70.0;
257
258         return std::make_pair(xmax, ymax);
259 }
260