git.sesse.net Git - pitch/blob - 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         std::pair<double, double> peak =
  97                 interpolate_peak(in[best_bin - 1],
  98                                  in[best_bin],
  99                                  in[best_bin + 1]);
 100
 101 #if 0
 102         printf("undertone strength: %+6.2f %+6.2f %+6.2f [%+6.2f] %+6.2f %+6.2f %+6.2f\n",
 103                 20.0 * log10(in[best_bin*4] / fft_length),
 104                 20.0 * log10(in[best_bin*3] / fft_length),
 105                 20.0 * log10(in[best_bin*2] / fft_length),
 106                 20.0 * log10(in[best_bin] / fft_length),
 107                 20.0 * log10(in[best_bin/2] / fft_length),
 108                 20.0 * log10(in[best_bin/3] / fft_length),
 109                 20.0 * log10(in[best_bin/4] / fft_length));
 110 #endif
 111
 112         // see if we might have hit an overtone (set a limit of 20dB)
 113         for (unsigned i = 6; i >= 1; --i) {
 114                 unsigned candidate_bin = best_bin / i;
 115                 if (best_bin == candidate_bin ||
 116                     candidate_bin < 5 ||
 117                     candidate_bin >= num_samples / 2 - 1) {
 118                         continue;
 119                 }
 120
 121                 // Check that we indeed have a peak in this area.
 122                 if (candidate_bin - 1 != best_bin &&
 123                     in[candidate_bin - 1] > in[candidate_bin] &&
 124                     in[candidate_bin - 1] > in[candidate_bin - 2]) {
 125                         --candidate_bin;
 126                 } else if (candidate_bin + 1 != best_bin &&
 127                            in[candidate_bin + 1] > in[candidate_bin] &&
 128                            in[candidate_bin + 1] > in[candidate_bin + 2]) {
 129                         ++candidate_bin;
 130                 } else if (in[candidate_bin] < in[candidate_bin - 1] ||
 131                            in[candidate_bin] < in[candidate_bin + 1]) {
 132                         continue;
 133                 }
 134
 135                 std::pair<double, double> candidate_peak =
 136                         interpolate_peak(in[candidate_bin - 1],
 137                                          in[candidate_bin],
 138                                          in[candidate_bin + 1]);
 139                 if (candidate_peak.second + 20.0f > peak.second) {
 140 #if 0
 141                         printf("Overtone of degree %u!\n", i);
 142 #endif
 143                         best_bin = candidate_bin;
 144                         peak = candidate_peak;
 145                         break;
 146                 }
 147         }
 148
 149         return std::make_pair(bin_to_freq(double(best_bin) + peak.first, num_samples), peak.second);
 150 }
 151
 152 // it's perhaps not ideal to _first_ find the peak and _then_ the harmonics --
 153 // ideally, one should find the set of all peaks and then determine the likely
 154 // base from that... something for later. :-)
 155 std::pair<double, double> PitchDetector::adjust_for_overtones(std::pair<double, double> base, double *in, unsigned num_samples)
 156 {
 157         double mu = base.first, var = 1.0 / (base.second * base.second);
 158
 159         //printf("Base at %.2f (amp=%5.2fdB)\n", base.first, base.second);
 160
 161         for (unsigned i = 2; i < 10; ++i) {
 162                 unsigned middle = unsigned(floor(freq_to_bin(base.first, num_samples) * i + 0.5));
 163                 unsigned lower = middle - (i+1)/2, upper = middle + (i+1)/2;
 164
 165                 if (lower < 1)
 166                         lower = 1;
 167                 if (upper >= num_samples)
 168                         upper = num_samples - 2;
 169
 170                 // printf("Searching in [%u,%u] = %f..%f\n", lower, upper, bin_to_freq(lower, num_samples), bin_to_freq(upper, num_samples));
 171
 172                 // search for a peak in this interval
 173                 double best_harmonics_freq = -1.0;
 174                 double best_harmonics_amp = -1.0;
 175                 for (unsigned j = lower; j <= upper; ++j) {
 176                         if (in[j] > in[j-1] && in[j] > in[j+1]) {
 177                                 std::pair<double, double> peak =
 178                                         interpolate_peak(in[j - 1],
 179                                                 in[j],
 180                                                 in[j + 1]);
 181
 182                                 if (peak.second > best_harmonics_amp) {
 183                                         best_harmonics_freq = bin_to_freq(j + peak.first, num_samples);
 184                                         best_harmonics_amp = peak.second;
 185                                 }
 186                         }
 187                 }
 188
 189                 if (best_harmonics_amp <= 0.0)
 190                         continue;
 191
 192                 //printf("Found overtone %u at %.2f (amp=%5.2fdB)\n", i, best_harmonics_freq,
 193                 //       best_harmonics_amp);
 194
 195                 double this_mu = best_harmonics_freq / double(i);
 196                 double this_var = 1.0 / (best_harmonics_amp * best_harmonics_amp);
 197
 198                 double k = var / (var + this_var);
 199                 mu = (1.0 - k) * mu + k * this_mu;
 200                 var *= (1.0 - k);
 201         }
 202         return std::make_pair(mu, base.second);
 203 }
 204
 205 double PitchDetector::bin_to_freq(double bin, unsigned num_samples)
 206 {
 207         return bin * sample_rate / double(num_samples);
 208 }
 209 double PitchDetector::freq_to_bin(double freq, unsigned num_samples)
 210 {
 211         return freq * double(num_samples) / double(sample_rate);
 212 }
 213
 214 /*
 215  * Given three bins, find the interpolated real peak based
 216  * on their magnitudes. To do this, we execute the following
 217  * plan:
 218  *
 219  * Fit a polynomial of the form ax^2 + bx + c = 0 to the data
 220  * we have. Maple readily yields our coefficients, assuming
 221  * that we have the values at x=-1, x=0 and x=1:
 222  *
 223  * > f := x -> a*x^2 + b*x + c;
 224  *
 225  *                               2
 226  *           f := x -> a x  + b x + c
 227  *
 228  * > cf := solve({ f(-1) = ym1, f(0) = y0, f(1) = y1 }, { a, b, c });
 229  *
 230  *                            y1    ym1       y1    ym1
 231  *        cf := {c = y0, b = ---- - ---, a = ---- + --- - y0}
 232  *                            2      2        2      2
 233  *
 234  * Now let's find the maximum point for the polynomial (it has to be
 235  * a maximum, since y0 is the greatest value):
 236  *
 237  * > xmax := solve(subs(cf, diff(f(x), x)) = 0, x);
 238  *
 239  *                                -y1 + ym1
 240  *                   xmax := -------------------
 241  *                           2 (y1 + ym1 - 2 y0)
 242  *
 243  * We could go further and insert {fmax,a,b,c} into the original
 244  * polynomial, but it just gets hairy. We calculate a, b and c separately
 245  * instead.
 246  *
 247  * http://www-ccrma.stanford.edu/~jos/parshl/Peak_Detection_Steps_3.html
 248  * claims that detection is almost twice as good when using dB scale instead
 249  * of linear scale for the input values, so we use that. (As a tiny bonus,
 250  * we get back dB scale from the function.)
 251  */
 252 std::pair<double, double> PitchDetector::interpolate_peak(double ym1, double y0, double y1)
 253 {
 254         ym1 = log10(ym1);
 255         y0 = log10(y0);
 256         y1 = log10(y1);
 257
 258 #if 0
 259         assert(y0 >= y1);
 260         assert(y0 >= ym1);
 261 #endif
 262
 263         double a = 0.5 * y1 + 0.5 * ym1 - y0;
 264         double b = 0.5 * y1 - 0.5 * ym1;
 265         double c = y0;
 266
 267         double xmax = (ym1 - y1) / (2.0 * (y1 + ym1 - 2.0 * y0));
 268         double ymax = 20.0 * (a * xmax * xmax + b * xmax + c) - 70.0;
 269
 270         return std::make_pair(xmax, ymax);
 271 }
 272