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 int get_dsp_fd();
  21 void read_chunk(int fd, double *in, unsigned num_samples);
  22 void apply_window(double *in, double *out, unsigned num_samples);
  23 std::pair<double, double> find_peak(double *in, unsigned num_samples);
  24 void find_peak_magnitudes(std::complex<double> *in, double *out, unsigned num_samples);
  25 double bin_to_freq(double bin, unsigned num_samples);
  26 std::string freq_to_tonename(double freq);
  27 std::pair<double, double> interpolate_peak(double ym1, double y0, double y1);
  28 void print_spectrogram(double freq, double amp);
  29 void write_sine(int dsp_fd, double freq, unsigned num_samples);
  30
  31 int main()
  32 {
  33         double *in, *in_windowed;
  34         std::complex<double> *out;
  35         double *bins;
  36         fftw_plan p;
  37
  38         // allocate memory
  39         in = reinterpret_cast<double *> (fftw_malloc(sizeof(double) * FFT_LENGTH / PAD_FACTOR));
  40         in_windowed = reinterpret_cast<double *> (fftw_malloc(sizeof(double) * FFT_LENGTH));
  41         out = reinterpret_cast<std::complex<double> *> (fftw_malloc(sizeof(std::complex<double>) * (FFT_LENGTH / 2 + 1)));
  42         bins = reinterpret_cast<double *> (fftw_malloc(sizeof(double) * FFT_LENGTH / 2 + 1));
  43
  44         memset(in, 0, sizeof(double) * FFT_LENGTH / PAD_FACTOR);
  45
  46         // init FFTW
  47         p = fftw_plan_dft_r2c_1d(FFT_LENGTH, in_windowed, reinterpret_cast<fftw_complex *> (out), FFTW_ESTIMATE);
  48
  49         int fd = get_dsp_fd();
  50         for ( ;; ) {
  51                 read_chunk(fd, in, FFT_LENGTH);
  52                 apply_window(in, in_windowed, FFT_LENGTH);
  53                 fftw_execute(p);
  54                 find_peak_magnitudes(out, bins, FFT_LENGTH);
  55                 std::pair<double, double> peak = find_peak(bins, FFT_LENGTH);
  56
  57                 if (peak.first < 50.0 || peak.second - log10(FFT_LENGTH) < 0.0) {
  58                         printf("............\n");
  59                 } else {
  60                         print_spectrogram(peak.first, peak.second - log10(FFT_LENGTH));
  61                 }
  62         }
  63 }
  64
  65 int get_dsp_fd()
  66 {
  67         int fd = open("/dev/dsp", O_RDWR);
  68         if (fd == -1) {
  69                 perror("/dev/dsp");
  70                 exit(1);
  71         }
  72
  73         ioctl(3, SNDCTL_DSP_RESET, 0);
  74
  75         int fmt = AFMT_S16_LE;   // FIXME
  76         ioctl(fd, SNDCTL_DSP_SETFMT, &fmt);
  77
  78         int chan = 1;
  79         ioctl(fd, SOUND_PCM_WRITE_CHANNELS, &chan);
  80
  81         int rate = 22050;
  82         ioctl(fd, SOUND_PCM_WRITE_RATE, &rate);
  83
  84         ioctl(3, SNDCTL_DSP_SYNC, 0);
  85
  86         return fd;
  87 }
  88
  89 #if 1
  90 void read_chunk(int fd, double *in, unsigned num_samples)
  91 {
  92         int ret;
  93         unsigned to_read = num_samples / PAD_FACTOR / OVERLAP;
  94         short buf[to_read];
  95
  96         memmove(in, in + to_read, (num_samples / PAD_FACTOR - to_read) * sizeof(double));
  97
  98         ret = read(fd, buf, to_read * sizeof(short));
  99         if (ret == 0) {
 100                 printf("EOF\n");
 101                 exit(0);
 102         }
 103
 104         if (ret != int(to_read * sizeof(short))) {
 105                 // blah
 106                 perror("read");
 107                 exit(1);
 108         }
 109
 110         for (unsigned i = 0; i < to_read; ++i)
 111                 in[i + (num_samples / PAD_FACTOR - to_read)] = double(buf[i]);
 112 }
 113 #else
 114 // make a pure 440hz sine for testing
 115 void read_chunk(int fd, double *in, unsigned num_samples)
 116 {
 117         static double theta = 0.0;
 118         for (unsigned i = 0; i < num_samples; ++i) {
 119                 in[i] = cos(theta);
 120                 theta += 2.0 * M_PI * 440.0 / double(SAMPLE_RATE);
 121         }
 122 }
 123 #endif
 124
 125 void write_sine(int dsp_fd, double freq, unsigned num_samples)
 126 {
 127         static double theta = 0.0;
 128         short buf[num_samples];
 129
 130         for (unsigned i = 0; i < num_samples; ++i) {
 131                 buf[i] = short(cos(theta) * 16384.0);
 132                 theta += 2.0 * M_PI * freq / double(SAMPLE_RATE);
 133         }
 134
 135         write(dsp_fd, buf, num_samples * sizeof(short));
 136 }
 137
 138 // Apply a standard Hamming window to our input data.
 139 void apply_window(double *in, double *out, unsigned num_samples)
 140 {
 141         static double *win_data;
 142         static unsigned win_len;
 143         static bool win_inited = false;
 144
 145         // Initialize the window for the first time
 146         if (!win_inited) {
 147                 win_len = num_samples / PAD_FACTOR;
 148                 win_data = new double[win_len];
 149
 150                 for (unsigned i = 0; i < win_len; ++i) {
 151                         win_data[i] = 0.54 - 0.46 * cos(2.0 * M_PI * double(i) / double(win_len - 1));
 152                 }
 153
 154                 win_inited = true;
 155         }
 156
 157         assert(win_len == num_samples / PAD_FACTOR);
 158
 159         for (unsigned i = 0; i < win_len; ++i) {
 160                 out[i] = in[i] * win_data[i];
 161         }
 162         for (unsigned i = win_len; i < num_samples; ++i) {
 163                 out[i] = 0.0;
 164         }
 165 }
 166
 167 void find_peak_magnitudes(std::complex<double> *in, double *out, unsigned num_samples)
 168 {
 169         for (unsigned i = 0; i < num_samples / 2 + 1; ++i)
 170                 out[i] = abs(in[i]);
 171 }
 172
 173 std::pair<double, double> find_peak(double *in, unsigned num_samples)
 174 {
 175         double best_peak = in[0];
 176         unsigned best_bin = 0;
 177
 178         for (unsigned i = 1; i < num_samples / 2 + 1; ++i) {
 179                 if (in[i] > best_peak) {
 180                         best_peak = in[i];
 181                         best_bin = i;
 182                 }
 183         }
 184
 185         if (best_bin == 0 || best_bin == num_samples / 2) {
 186                 return std::make_pair(-1.0, 0.0);
 187         }
 188
 189 #if 0
 190         printf("undertone strength: %+4.2f %+4.2f %+4.2f [%+4.2f] %+4.2f %+4.2f %+4.2f\n",
 191                 20.0 * log10(in[best_bin*4] / FFT_LENGTH),
 192                 20.0 * log10(in[best_bin*3] / FFT_LENGTH),
 193                 20.0 * log10(in[best_bin*2] / FFT_LENGTH),
 194                 20.0 * log10(in[best_bin] / FFT_LENGTH),
 195                 20.0 * log10(in[best_bin/2] / FFT_LENGTH),
 196                 20.0 * log10(in[best_bin/3] / FFT_LENGTH),
 197                 20.0 * log10(in[best_bin/4] / FFT_LENGTH));
 198 #endif
 199
 200         // see if we might have hit an overtone (set a limit of 5dB)
 201         for (unsigned i = 4; i >= 1; --i) {
 202                 if (best_bin != best_bin / i &&
 203                     20.0 * log10(in[best_bin] / in[best_bin / i]) < 5.0f) {
 204 #if 0
 205                         printf("Overtone of degree %u!\n", i);
 206 #endif
 207                         best_bin /= i;
 208
 209                         // consider sliding one bin up or down
 210                         if (best_bin > 0 && in[best_bin - 1] > in[best_bin] && in[best_bin - 1] > in[best_bin - 2]) {
 211                                 --best_bin;
 212                         } else if (best_bin < num_samples / 2 && in[best_bin + 1] > in[best_bin] && in[best_bin + 1] > in[best_bin + 2]) {
 213                                 ++best_bin;
 214                         }
 215
 216                         break;
 217                 }
 218         }
 219
 220         std::pair<double, double> peak =
 221                 interpolate_peak(in[best_bin - 1],
 222                                  in[best_bin],
 223                                  in[best_bin + 1]);
 224
 225         return std::make_pair(bin_to_freq(double(best_bin) + peak.first, num_samples), peak.second);
 226 }
 227
 228 double bin_to_freq(double bin, unsigned num_samples)
 229 {
 230         return bin * SAMPLE_RATE / double(num_samples);
 231 }
 232
 233 /*
 234  * Given three bins, find the interpolated real peak based
 235  * on their magnitudes. To do this, we execute the following
 236  * plan:
 237  *
 238  * Fit a polynomial of the form ax^2 + bx + c = 0 to the data
 239  * we have. Maple readily yields our coefficients, assuming
 240  * that we have the values at x=-1, x=0 and x=1:
 241  *
 242  * > f := x -> a*x^2 + b*x + c;
 243  *
 244  *                               2
 245  *           f := x -> a x  + b x + c
 246  *
 247  * > cf := solve({ f(-1) = ym1, f(0) = y0, f(1) = y1 }, { a, b, c });
 248  *
 249  *                            y1    ym1       y1    ym1
 250  *        cf := {c = y0, b = ---- - ---, a = ---- + --- - y0}
 251  *                            2      2        2      2
 252  *
 253  * Now let's find the maximum point for the polynomial (it has to be
 254  * a maximum, since y0 is the greatest value):
 255  *
 256  * > xmax := solve(subs(cf, diff(f(x), x)) = 0, x);
 257  *
 258  *                                -y1 + ym1
 259  *                   xmax := -------------------
 260  *                           2 (y1 + ym1 - 2 y0)
 261  *
 262  * We could go further and insert {fmax,a,b,c} into the original
 263  * polynomial, but it just gets hairy. We calculate a, b and c separately
 264  * instead.
 265  *
 266  * http://www-ccrma.stanford.edu/~jos/parshl/Peak_Detection_Steps_3.html
 267  * claims that detection is almost twice as good when using dB scale instead
 268  * of linear scale for the input values, so we use that. (As a tiny bonus,
 269  * we get back dB scale from the function.)
 270  */
 271 std::pair<double, double> interpolate_peak(double ym1, double y0, double y1)
 272 {
 273         ym1 = log10(ym1);
 274         y0 = log10(y0);
 275         y1 = log10(y1);
 276
 277 #if 0
 278         assert(y0 >= y1);
 279         assert(y0 >= ym1);
 280 #endif
 281
 282         double a = 0.5 * y1 + 0.5 * ym1 - y0;
 283         double b = 0.5 * y1 - 0.5 * ym1;
 284         double c = y0;
 285
 286         double xmax = (ym1 - y1) / (2.0 * (y1 + ym1 - 2.0 * y0));
 287         double ymax = 20.0 * (a * xmax * xmax + b * xmax + c) - 90.0;
 288
 289         return std::make_pair(xmax, ymax);
 290 }
 291
 292 std::string freq_to_tonename(double freq)
 293 {
 294         std::string notenames[] = { "C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B" };
 295         double half_notes_away = 12.0 * log2(freq / 440.0) - 3.0;
 296         int hnai = int(floor(half_notes_away + 0.5));
 297         int octave = (hnai + 48) / 12;
 298
 299         char buf[256];
 300         sprintf(buf, "%s%d + %.2f [%d]", notenames[((hnai % 12) + 12) % 12].c_str(), octave, half_notes_away - hnai, hnai);
 301         return buf;
 302 }
 303
 304 void print_spectrogram(double freq, double amp)
 305 {
 306         std::string notenames[] = { "C", "C#", "D", "D#", "E", "F", "F#", "G", "G#", "A", "A#", "B" };
 307         double half_notes_away = 12.0 * log2(freq / 440.0) - 3.0;
 308         int hnai = int(floor(half_notes_away + 0.5));
 309         int octave = (hnai + 48) / 12;
 310
 311         for (int i = 0; i < 12; ++i)
 312                 if (i == ((hnai % 12) + 12) % 12)
 313                         printf("#");
 314                 else
 315                         printf(".");
 316
 317         printf(" (%-2s%d %+.2f, %5.2fHz) [%5.2fdB]  [", notenames[((hnai % 12) + 12) % 12].c_str(), octave, half_notes_away - hnai,
 318                 freq, amp);
 319
 320         double off = half_notes_away - hnai;
 321         for (int i = -10; i <= 10; ++i) {
 322                 if (off >= (i-0.5) * 0.05 && off < (i+0.5) * 0.05) {
 323                         printf("#");
 324                 } else {
 325                         if (i == 0) {
 326                                 printf("|");
 327                         } else {
 328                                 printf("-");
 329                         }
 330                 }
 331         }
 332         printf("]\n");
 333
 334 }