git.sesse.net Git - pitch/blob - pitch.cpp

   1 #include <stdio.h>
   2 #include <stdlib.h>
   3 #include <unistd.h>
   4 #include <fcntl.h>
   5 #include <complex>
   6 #include <cassert>
   7 #include <algorithm>
   8 #include <fftw3.h>
   9 #include <sys/ioctl.h>
  10 #include <linux/soundcard.h>
  11
  12 #define SAMPLE_RATE     22050
  13 #define FFT_LENGTH      4096     /* in samples */
  14 #define PAD_FACTOR      2        /* 1/pf of the FFT samples are real samples, the rest are padding */
  15 #define OVERLAP         4        /* 1/ol samples will be replaced in the buffer every frame. Should be
  16                                   * a multiple of 2 for the Hamming window (see
  17                                   * http://www-ccrma.stanford.edu/~jos/parshl/Choice_Hop_Size.html).
  18                                   */
  19
  20 #define EQUAL_TEMPERAMENT     0
  21 #define WELL_TEMPERED_GUITAR  1
  22
  23 #define TUNING WELL_TEMPERED_GUITAR
  24
  25 int get_dsp_fd();
  26 void read_chunk(int fd, double *in, unsigned num_samples);
  27 void apply_window(double *in, double *out, unsigned num_samples);
  28 std::pair<double, double> find_peak(double *in, unsigned num_samples);
  29 void find_peak_magnitudes(std::complex<double> *in, double *out, unsigned num_samples);
  30 std::pair<double, double> adjust_for_overtones(std::pair<double, double> base, double *in, unsigned num_samples);
  31 double bin_to_freq(double bin, unsigned num_samples);
  32 double freq_to_bin(double freq, unsigned num_samples);
  33 std::string freq_to_tonename(double freq);
  34 std::pair<double, double> interpolate_peak(double ym1, double y0, double y1);
  35 void print_spectrogram(double freq, double amp);
  36 void write_sine(int dsp_fd, double freq, unsigned num_samples);
  37
  38 int main()
  39 {
  40         double *in, *in_windowed;
  41         std::complex<double> *out;
  42         double *bins;
  43         fftw_plan p;
  44
  45         // allocate memory
  46         in = reinterpret_cast<double *> (fftw_malloc(sizeof(double) * FFT_LENGTH / PAD_FACTOR));
  47         in_windowed = reinterpret_cast<double *> (fftw_malloc(sizeof(double) * FFT_LENGTH));
  48         out = reinterpret_cast<std::complex<double> *> (fftw_malloc(sizeof(std::complex<double>) * (FFT_LENGTH / 2 + 1)));
  49         bins = reinterpret_cast<double *> (fftw_malloc(sizeof(double) * FFT_LENGTH / 2 + 1));
  50
  51         memset(in, 0, sizeof(double) * FFT_LENGTH / PAD_FACTOR);
  52
  53         // init FFTW
  54         p = fftw_plan_dft_r2c_1d(FFT_LENGTH, in_windowed, reinterpret_cast<fftw_complex *> (out), FFTW_ESTIMATE);
  55
  56         int fd = get_dsp_fd();
  57         for ( ;; ) {
  58                 read_chunk(fd, in, FFT_LENGTH);
  59                 apply_window(in, in_windowed, FFT_LENGTH);
  60                 fftw_execute(p);
  61                 find_peak_magnitudes(out, bins, FFT_LENGTH);
  62                 std::pair<double, double> peak = find_peak(bins, FFT_LENGTH);
  63                 peak = adjust_for_overtones(peak, bins, FFT_LENGTH);
  64
  65                 if (peak.first < 50.0 || peak.second - log10(FFT_LENGTH) < 0.0) {
  66 #if TUNING == WELL_TEMPERED_GUITAR
  67                         printf("......\n");
  68 #else
  69                         printf("............\n");
  70 #endif
  71                 } else {
  72                         print_spectrogram(peak.first, peak.second - log10(FFT_LENGTH));
  73                 }
  74         }
  75 }
  76
  77 int get_dsp_fd()
  78 {
  79         int fd = open("/dev/dsp", O_RDWR);
  80         if (fd == -1) {
  81                 perror("/dev/dsp");
  82                 exit(1);
  83         }
  84
  85         ioctl(3, SNDCTL_DSP_RESET, 0);
  86
  87         int fmt = AFMT_S16_LE;   // FIXME
  88         ioctl(fd, SNDCTL_DSP_SETFMT, &fmt);
  89
  90         int chan = 1;
  91         ioctl(fd, SOUND_PCM_WRITE_CHANNELS, &chan);
  92
  93         int rate = SAMPLE_RATE;
  94         ioctl(fd, SOUND_PCM_WRITE_RATE, &rate);
  95
  96         int fragments = 0x00020002;
  97         ioctl(fd, SNDCTL_DSP_SETFRAGMENT, &fragments);
  98
  99         ioctl(3, SNDCTL_DSP_SYNC, 0);
 100
 101         return fd;
 102 }
 103
 104 #if 1
 105 void read_chunk(int fd, double *in, unsigned num_samples)
 106 {
 107         int ret;
 108         unsigned to_read = num_samples / PAD_FACTOR / OVERLAP;
 109         short buf[to_read];
 110
 111         memmove(in, in + to_read, (num_samples / PAD_FACTOR - to_read) * sizeof(double));
 112
 113         ret = read(fd, buf, to_read * sizeof(short));
 114         if (ret == 0) {
 115                 printf("EOF\n");
 116                 exit(0);
 117         }
 118
 119         if (ret != int(to_read * sizeof(short))) {
 120                 // blah
 121                 perror("read");
 122                 exit(1);
 123         }
 124
 125         for (unsigned i = 0; i < to_read; ++i)
 126                 in[i + (num_samples / PAD_FACTOR - to_read)] = double(buf[i]);
 127 }
 128 #else
 129 // make a pure 440hz sine for testing
 130 void read_chunk(int fd, double *in, unsigned num_samples)
 131 {
 132         static double theta = 0.0;
 133         for (unsigned i = 0; i < num_samples; ++i) {
 134                 in[i] = cos(theta);
 135                 theta += 2.0 * M_PI * 440.0 / double(SAMPLE_RATE);
 136         }
 137 }
 138 #endif
 139
 140 void write_sine(int dsp_fd, double freq, unsigned num_samples)
 141 {
 142         static double theta = 0.0;
 143         short buf[num_samples];
 144
 145         for (unsigned i = 0; i < num_samples; ++i) {
 146                 buf[i] = short(cos(theta) * 16384.0);
 147                 theta += 2.0 * M_PI * freq / double(SAMPLE_RATE);
 148         }
 149
 150         write(dsp_fd, buf, num_samples * sizeof(short));
 151 }
 152
 153 // Apply a standard Hamming window to our input data.
 154 void apply_window(double *in, double *out, unsigned num_samples)
 155 {
 156         static double *win_data;
 157         static unsigned win_len;
 158         static bool win_inited = false;
 159
 160         // Initialize the window for the first time
 161         if (!win_inited) {
 162                 win_len = num_samples / PAD_FACTOR;
 163                 win_data = new double[win_len];
 164
 165                 for (unsigned i = 0; i < win_len; ++i) {
 166                         win_data[i] = 0.54 - 0.46 * cos(2.0 * M_PI * double(i) / double(win_len - 1));
 167                 }
 168
 169                 win_inited = true;
 170         }
 171
 172         assert(win_len == num_samples / PAD_FACTOR);
 173
 174         for (unsigned i = 0; i < win_len; ++i) {
 175                 out[i] = in[i] * win_data[i];
 176         }
 177         for (unsigned i = win_len; i < num_samples; ++i) {
 178                 out[i] = 0.0;
 179         }
 180 }
 181
 182 void find_peak_magnitudes(std::complex<double> *in, double *out, unsigned num_samples)
 183 {
 184         for (unsigned i = 0; i < num_samples / 2 + 1; ++i)
 185                 out[i] = abs(in[i]);
 186 }
 187
 188 std::pair<double, double> find_peak(double *in, unsigned num_samples)
 189 {
 190         double best_peak = in[0];
 191         unsigned best_bin = 0;
 192
 193         for (unsigned i = 1; i < num_samples / 2 + 1; ++i) {
 194                 if (in[i] > best_peak) {
 195                         best_peak = in[i];
 196                         best_bin = i;
 197                 }
 198         }
 199
 200         if (best_bin == 0 || best_bin == num_samples / 2) {
 201                 return std::make_pair(-1.0, 0.0);
 202         }
 203
 204 #if 0
 205         printf("undertone strength: %+4.2f %+4.2f %+4.2f [%+4.2f] %+4.2f %+4.2f %+4.2f\n",
 206                 20.0 * log10(in[best_bin*4] / FFT_LENGTH),
 207                 20.0 * log10(in[best_bin*3] / FFT_LENGTH),
 208                 20.0 * log10(in[best_bin*2] / FFT_LENGTH),
 209                 20.0 * log10(in[best_bin] / FFT_LENGTH),
 210                 20.0 * log10(in[best_bin/2] / FFT_LENGTH),
 211                 20.0 * log10(in[best_bin/3] / FFT_LENGTH),
 212                 20.0 * log10(in[best_bin/4] / FFT_LENGTH));
 213 #endif
 214
 215         // see if we might have hit an overtone (set a limit of 5dB)
 216         for (unsigned i = 4; i >= 1; --i) {
 217                 if (best_bin != best_bin / i &&
 218                     20.0 * log10(in[best_bin] / in[best_bin / i]) < 5.0f) {
 219 #if 0
 220                         printf("Overtone of degree %u!\n", i);
 221 #endif
 222                         best_bin /= i;
 223
 224                         // consider sliding one bin up or down
 225                         if (best_bin > 0 && in[best_bin - 1] > in[best_bin] && in[best_bin - 1] > in[best_bin - 2]) {
 226                                 --best_bin;
 227                         } else if (best_bin < num_samples / 2 && in[best_bin + 1] > in[best_bin] && in[best_bin + 1] > in[best_bin + 2]) {
 228                                 ++best_bin;
 229                         }
 230
 231                         break;
 232                 }
 233         }
 234
 235         std::pair<double, double> peak =
 236                 interpolate_peak(in[best_bin - 1],
 237                                  in[best_bin],
 238                                  in[best_bin + 1]);
 239
 240         return std::make_pair(bin_to_freq(double(best_bin) + peak.first, num_samples), peak.second);
 241 }
 242
 243 // it's perhaps not ideal to _first_ find the peak and _then_ the harmonics --
 244 // ideally, one should find the set of all peaks and then determine the likely
 245 // base from that... something for later. :-)
 246 std::pair<double, double> adjust_for_overtones(std::pair<double, double> base, double *in, unsigned num_samples)
 247 {
 248         double mu = base.first, var = 1.0 / (base.second * base.second);
 249         printf("mu=%f, var=%f\n", mu, var);
 250
 251         for (unsigned i = 2; i < 10; ++i) {
 252                 unsigned middle = unsigned(floor(freq_to_bin(base.first, num_samples) * i + 0.5));
 253                 unsigned lower = middle - (i+1)/2, upper = middle + (i+1)/2;
 254
 255                 if (upper >= num_samples)
 256                         upper = num_samples - 2;
 257
 258                 // printf("Searching in [%u,%u] = %f..%f\n", lower, upper, bin_to_freq(lower, num_samples), bin_to_freq(upper, num_samples));
 259
 260                 // search for a peak in this interval
 261                 double best_harmonics_freq = -1.0;
 262                 double best_harmonics_amp = -1.0;
 263                 for (unsigned j = lower; j <= upper; ++j) {
 264                         if (in[j] > in[j-1] && in[j] > in[j+1]) {
 265                                 std::pair<double, double> peak =
 266                                         interpolate_peak(in[j - 1],
 267                                                 in[j],
 268                                                 in[j + 1]);
 269
 270                                 if (peak.second > best_harmonics_amp) {
 271                                         best_harmonics_freq = bin_to_freq(j + peak.first, num_samples);
 272                                         best_harmonics_amp = peak.second;
 273                                 }
 274                         }
 275                 }
 276
 277                 if (best_harmonics_amp <= 0.0)
 278                         continue;
 279
 280                 //printf("Found overtone %u at %.2f (amp=%5.2fdB)\n", i, best_harmonics_freq,
 281                 //       best_harmonics_amp);
 282
 283                 double this_mu = best_harmonics_freq / double(i);
 284                 double this_var = 1.0 / (best_harmonics_amp * best_harmonics_amp);
 285
 286                 double k = var / (var + this_var);
 287                 mu = (1.0 - k) * mu + k * this_mu;
 288                 var *= (1.0 - k);
 289
 290                 printf("mu=%f, var=%f\n", mu, var);
 291         }
 292         return std::make_pair(mu, base.second);
 293 }
 294
 295 double bin_to_freq(double bin, unsigned num_samples)
 296 {
 297         return bin * SAMPLE_RATE / double(num_samples);
 298 }
 299 double freq_to_bin(double freq, unsigned num_samples)
 300 {
 301         return freq * double(num_samples) / double(SAMPLE_RATE);
 302 }
 303
 304 /*
 305  * Given three bins, find the interpolated real peak based
 306  * on their magnitudes. To do this, we execute the following
 307  * plan:
 308  *
 309  * Fit a polynomial of the form ax^2 + bx + c = 0 to the data
 310  * we have. Maple readily yields our coefficients, assuming
 311  * that we have the values at x=-1, x=0 and x=1:
 312  *
 313  * > f := x -> a*x^2 + b*x + c;
 314  *
 315  *                               2
 316  *           f := x -> a x  + b x + c
 317  *
 318  * > cf := solve({ f(-1) = ym1, f(0) = y0, f(1) = y1 }, { a, b, c });
 319  *
 320  *                            y1    ym1       y1    ym1
 321  *        cf := {c = y0, b = ---- - ---, a = ---- + --- - y0}
 322  *                            2      2        2      2
 323  *
 324  * Now let's find the maximum point for the polynomial (it has to be
 325  * a maximum, since y0 is the greatest value):
 326  *
 327  * > xmax := solve(subs(cf, diff(f(x), x)) = 0, x);
 328  *
 329  *                                -y1 + ym1
 330  *                   xmax := -------------------
 331  *                           2 (y1 + ym1 - 2 y0)
 332  *
 333  * We could go further and insert {fmax,a,b,c} into the original
 334  * polynomial, but it just gets hairy. We calculate a, b and c separately
 335  * instead.
 336  *
 337  * http://www-ccrma.stanford.edu/~jos/parshl/Peak_Detection_Steps_3.html
 338  * claims that detection is almost twice as good when using dB scale instead
 339  * of linear scale for the input values, so we use that. (As a tiny bonus,
 340  * we get back dB scale from the function.)
 341  */
 342 std::pair<double, double> interpolate_peak(double ym1, double y0, double y1)
 343 {
 344         ym1 = log10(ym1);
 345         y0 = log10(y0);
 346         y1 = log10(y1);
 347
 348 #if 0
 349         assert(y0 >= y1);
 350         assert(y0 >= ym1);
 351 #endif
 352
 353         double a = 0.5 * y1 + 0.5 * ym1 - y0;
 354         double b = 0.5 * y1 - 0.5 * ym1;
 355         double c = y0;
 356
 357         double xmax = (ym1 - y1) / (2.0 * (y1 + ym1 - 2.0 * y0));
 358         double ymax = 20.0 * (a * xmax * xmax + b * xmax + c) - 90.0;
 359
 360         return std::make_pair(xmax, ymax);
 361 }
 362
 363 std::string freq_to_tonename(double freq)
 364 {
 365         std::string notenames[] = { "C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B" };
 366         double half_notes_away = 12.0 * log2(freq / 440.0) - 3.0;
 367         int hnai = int(floor(half_notes_away + 0.5));
 368         int octave = (hnai + 48) / 12;
 369
 370         char buf[256];
 371         sprintf(buf, "%s%d + %.2f [%d]", notenames[((hnai % 12) + 12) % 12].c_str(), octave, half_notes_away - hnai, hnai);
 372         return buf;
 373 }
 374
 375 #if TUNING == EQUAL_TEMPERAMENT
 376 void print_spectrogram(double freq, double amp)
 377 {
 378         std::string notenames[] = { "C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B" };
 379         double half_notes_away = 12.0 * log2(freq / 440.0) - 3.0;
 380         int hnai = int(floor(half_notes_away + 0.5));
 381         int octave = (hnai + 48) / 12;
 382
 383         for (int i = 0; i < 12; ++i)
 384                 if (i == ((hnai % 12) + 12) % 12)
 385                         printf("#");
 386                 else
 387                         printf(".");
 388
 389         printf(" (%-2s%d %+.2f, %5.2fHz) [%5.2fdB]  [", notenames[((hnai % 12) + 12) % 12].c_str(), octave, half_notes_away - hnai,
 390                 freq, amp);
 391
 392         double off = half_notes_away - hnai;
 393         for (int i = -10; i <= 10; ++i) {
 394                 if (off >= (i-0.5) * 0.05 && off < (i+0.5) * 0.05) {
 395                         printf("#");
 396                 } else {
 397                         if (i == 0) {
 398                                 printf("|");
 399                         } else {
 400                                 printf("-");
 401                         }
 402                 }
 403         }
 404         printf("]\n");
 405
 406 }
 407 #else
 408 struct note {
 409         char notename[16];
 410         double freq;
 411 };
 412 static note notes[] = {
 413         { "E-3", 110.0 * (3.0/4.0) },
 414         { "A-3", 110.0 },
 415         { "D-4", 110.0 * (4.0/3.0) },
 416         { "G-4", 110.0 * (4.0/3.0)*(4.0/3.0) },
 417         { "B-4", 440.0 * (3.0/4.0)*(3.0/4.0) },
 418         { "E-5", 440.0 * (3.0/4.0) }
 419 };
 420
 421 void print_spectrogram(double freq, double amp)
 422 {
 423         double best_away = 999999999.9;
 424         unsigned best_away_ind = 0;
 425
 426         for (unsigned i = 0; i < sizeof(notes)/sizeof(note); ++i) {
 427                 double half_notes_away = 12.0 * log2(freq / notes[i].freq);
 428                 if (fabs(half_notes_away) < fabs(best_away)) {
 429                         best_away = half_notes_away;
 430                         best_away_ind = i;
 431                 }
 432         }
 433
 434         for (unsigned i = 0; i < sizeof(notes)/sizeof(note); ++i)
 435                 if (i == best_away_ind)
 436                         printf("#");
 437                 else
 438                         printf(".");
 439
 440         printf(" (%s %+.2f, %5.2fHz) [%5.2fdB]  [", notes[best_away_ind].notename, best_away, freq, amp);
 441
 442         for (int i = -10; i <= 10; ++i) {
 443                 if (best_away >= (i-0.5) * 0.05 && best_away < (i+0.5) * 0.05) {
 444                         printf("#");
 445                 } else {
 446                         if (i == 0) {
 447                                 printf("|");
 448                         } else {
 449                                 printf("-");
 450                         }
 451                 }
 452         }
 453         printf("]\n");
 454 }
 455 #endif