]> git.sesse.net Git - pitch/blobdiff - pitchdetector.cpp
Split the pitch logic into its own class.
[pitch] / pitchdetector.cpp
diff --git a/pitchdetector.cpp b/pitchdetector.cpp
new file mode 100644 (file)
index 0000000..71e5875
--- /dev/null
@@ -0,0 +1,251 @@
+#include <stdio.h>
+#include <stdlib.h>
+#include <unistd.h>
+#include <fcntl.h>
+#include <complex>
+#include <cassert>
+#include <algorithm>
+#include <fftw3.h>
+#include "pitchdetector.h"
+
+PitchDetector::PitchDetector(unsigned sample_rate, unsigned fft_length, unsigned pad_factor, unsigned overlap)
+       : sample_rate(sample_rate), fft_length(fft_length), pad_factor(pad_factor), overlap(overlap)
+{
+       in = reinterpret_cast<double *> (fftw_malloc(sizeof(double) * fft_length / pad_factor));
+       in_windowed = reinterpret_cast<double *> (fftw_malloc(sizeof(double) * fft_length));
+       out = reinterpret_cast<std::complex<double> *> (fftw_malloc(sizeof(std::complex<double>) * (fft_length / 2 + 1)));
+       bins = reinterpret_cast<double *> (fftw_malloc(sizeof(double) * (fft_length / 2 + 1)));
+
+       memset(in, 0, sizeof(double) * fft_length / pad_factor);
+
+       plan = fftw_plan_dft_r2c_1d(fft_length, in_windowed, reinterpret_cast<fftw_complex *> (out), FFTW_ESTIMATE);
+       
+       // Initialize the Hamming window
+       window_data = new double[fft_length / pad_factor];
+       for (unsigned i = 0; i < fft_length / pad_factor; ++i) {
+               window_data[i] = 0.54 - 0.46 * cos(2.0 * M_PI * double(i) / double(fft_length/pad_factor - 1));
+       }
+}
+
+PitchDetector::~PitchDetector()
+{
+       fftw_free(in);
+       fftw_free(in_windowed);
+       fftw_free(out);
+       fftw_free(bins);
+}
+
+std::pair<double, double> PitchDetector::detect_pitch(short *buf)
+{
+       unsigned buf_len = fft_length / pad_factor / overlap;
+       memmove(in, in + buf_len, (fft_length - buf_len) * sizeof(double));
+       
+       for (unsigned i = 0; i < buf_len; ++i)
+               in[i + (fft_length / pad_factor - buf_len)] = double(buf[i]);
+
+       apply_window(in, in_windowed, fft_length);
+       fftw_execute(plan);
+       find_peak_magnitudes(out, bins, fft_length);
+       std::pair<double, double> peak = find_peak(bins, fft_length);
+       if (peak.first > 0.0)
+               peak = adjust_for_overtones(peak, bins, fft_length);
+
+       return peak;
+}
+
+// Apply a standard Hamming window to our input data.
+void PitchDetector::apply_window(double *in, double *out, unsigned num_samples)
+{
+       for (unsigned i = 0; i < num_samples / pad_factor; ++i) {
+               out[i] = in[i] * window_data[i];
+       }
+       for (unsigned i = num_samples / pad_factor; i < num_samples; ++i) {
+               out[i] = 0.0;
+       }
+}
+
+void PitchDetector::find_peak_magnitudes(std::complex<double> *in, double *out, unsigned num_samples)
+{
+       for (unsigned i = 0; i < num_samples / 2 + 1; ++i)
+               out[i] = abs(in[i]);
+}
+
+std::pair<double, double> PitchDetector::find_peak(double *in, unsigned num_samples)
+{
+       double best_peak = in[5];
+       unsigned best_bin = 5;
+
+       for (unsigned i = 6; i < num_samples / 2 + 1; ++i) {
+               if (in[i] > best_peak) {
+                       best_peak = in[i];
+                       best_bin = i;
+               }
+       }
+       
+       if (best_bin == 0 || best_bin == num_samples / 2) {
+               return std::make_pair(-1.0, 0.0);
+       }
+
+#if 0
+       printf("undertone strength: %+4.2f %+4.2f %+4.2f [%+4.2f] %+4.2f %+4.2f %+4.2f\n",
+               20.0 * log10(in[best_bin*4] / fft_length),
+               20.0 * log10(in[best_bin*3] / fft_length),
+               20.0 * log10(in[best_bin*2] / fft_length),
+               20.0 * log10(in[best_bin] / fft_length),
+               20.0 * log10(in[best_bin/2] / fft_length),
+               20.0 * log10(in[best_bin/3] / fft_length),
+               20.0 * log10(in[best_bin/4] / fft_length));
+#endif
+
+       // see if we might have hit an overtone (set a limit of 5dB)
+       for (unsigned i = 4; i >= 1; --i) {
+               if (best_bin != best_bin / i &&
+                   20.0 * log10(in[best_bin] / in[best_bin / i]) < 5.0f) {
+#if 0
+                       printf("Overtone of degree %u!\n", i);
+#endif
+                       best_bin /= i;
+
+                       // consider sliding one bin up or down
+                       if (best_bin > 1 && in[best_bin - 1] > in[best_bin] && in[best_bin - 1] > in[best_bin - 2]) {
+                               --best_bin;
+                       } else if (best_bin < num_samples / 2 - 1 && in[best_bin + 1] > in[best_bin] && in[best_bin + 1] > in[best_bin + 2]) {
+                               ++best_bin;
+                       }
+                          
+                       break;
+               }
+       }
+
+       if (best_bin == 0 || best_bin == num_samples / 2) {
+               return std::make_pair(-1.0, 0.0);
+       }
+       std::pair<double, double> peak = 
+               interpolate_peak(in[best_bin - 1],
+                                in[best_bin],
+                                in[best_bin + 1]);
+               
+       return std::make_pair(bin_to_freq(double(best_bin) + peak.first, num_samples), peak.second);
+}
+
+// it's perhaps not ideal to _first_ find the peak and _then_ the harmonics --
+// ideally, one should find the set of all peaks and then determine the likely
+// base from that... something for later. :-)
+std::pair<double, double> PitchDetector::adjust_for_overtones(std::pair<double, double> base, double *in, unsigned num_samples)
+{
+       double mu = base.first, var = 1.0 / (base.second * base.second);
+               
+       //printf("Base at %.2f (amp=%5.2fdB)\n", base.first, base.second);
+
+       for (unsigned i = 2; i < 10; ++i) {
+               unsigned middle = unsigned(floor(freq_to_bin(base.first, num_samples) * i + 0.5));
+               unsigned lower = middle - (i+1)/2, upper = middle + (i+1)/2;
+
+               if (lower < 1)
+                       lower = 1;
+               if (upper >= num_samples)
+                       upper = num_samples - 2;
+
+               // printf("Searching in [%u,%u] = %f..%f\n", lower, upper, bin_to_freq(lower, num_samples), bin_to_freq(upper, num_samples));
+
+               // search for a peak in this interval
+               double best_harmonics_freq = -1.0;
+               double best_harmonics_amp = -1.0;
+               for (unsigned j = lower; j <= upper; ++j) {
+                       if (in[j] > in[j-1] && in[j] > in[j+1]) {
+                               std::pair<double, double> peak =
+                                       interpolate_peak(in[j - 1],
+                                               in[j],
+                                               in[j + 1]);
+                               
+                               if (peak.second > best_harmonics_amp) {
+                                       best_harmonics_freq = bin_to_freq(j + peak.first, num_samples);
+                                       best_harmonics_amp = peak.second;
+                               }
+                       }
+               }
+
+               if (best_harmonics_amp <= 0.0)
+                       continue;
+
+               //printf("Found overtone %u at %.2f (amp=%5.2fdB)\n", i, best_harmonics_freq,
+               //       best_harmonics_amp);
+
+               double this_mu = best_harmonics_freq / double(i);
+               double this_var = 1.0 / (best_harmonics_amp * best_harmonics_amp);
+
+               double k = var / (var + this_var);
+               mu = (1.0 - k) * mu + k * this_mu;
+               var *= (1.0 - k);
+       }
+       return std::make_pair(mu, base.second);
+}
+
+double PitchDetector::bin_to_freq(double bin, unsigned num_samples)
+{
+       return bin * sample_rate / double(num_samples);
+}
+double PitchDetector::freq_to_bin(double freq, unsigned num_samples)
+{
+       return freq * double(num_samples) / double(sample_rate);
+}
+
+/*
+ * Given three bins, find the interpolated real peak based
+ * on their magnitudes. To do this, we execute the following
+ * plan:
+ * 
+ * Fit a polynomial of the form ax^2 + bx + c = 0 to the data
+ * we have. Maple readily yields our coefficients, assuming
+ * that we have the values at x=-1, x=0 and x=1:
+ *
+ * > f := x -> a*x^2 + b*x + c;                                
+ * 
+ *                               2
+ *           f := x -> a x  + b x + c
+ * 
+ * > cf := solve({ f(-1) = ym1, f(0) = y0, f(1) = y1 }, { a, b, c });
+ * 
+ *                            y1    ym1       y1    ym1
+ *        cf := {c = y0, b = ---- - ---, a = ---- + --- - y0}
+ *                            2      2        2      2
+ *
+ * Now let's find the maximum point for the polynomial (it has to be
+ * a maximum, since y0 is the greatest value):
+ *
+ * > xmax := solve(subs(cf, diff(f(x), x)) = 0, x);
+ * 
+ *                                -y1 + ym1
+ *                   xmax := -------------------
+ *                           2 (y1 + ym1 - 2 y0)
+ *
+ * We could go further and insert {fmax,a,b,c} into the original
+ * polynomial, but it just gets hairy. We calculate a, b and c separately
+ * instead.
+ *
+ * http://www-ccrma.stanford.edu/~jos/parshl/Peak_Detection_Steps_3.html
+ * claims that detection is almost twice as good when using dB scale instead
+ * of linear scale for the input values, so we use that. (As a tiny bonus,
+ * we get back dB scale from the function.)
+ */
+std::pair<double, double> PitchDetector::interpolate_peak(double ym1, double y0, double y1)
+{
+       ym1 = log10(ym1);
+       y0 = log10(y0);
+       y1 = log10(y1);
+
+#if 0
+       assert(y0 >= y1);
+       assert(y0 >= ym1);
+#endif
+       
+       double a = 0.5 * y1 + 0.5 * ym1 - y0;
+       double b = 0.5 * y1 - 0.5 * ym1;
+       double c = y0;
+       
+       double xmax = (ym1 - y1) / (2.0 * (y1 + ym1 - 2.0 * y0));
+       double ymax = 20.0 * (a * xmax * xmax + b * xmax + c) - 90.0;
+
+       return std::make_pair(xmax, ymax);
+}
+