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