compute per-band pitch

1 files changed, 39 insertions, 23 deletions

diff --git a/src/denoise.c b/src/denoise.c
index 10f5e33..4d14383 100644
--- a/src/denoise.c
+++ b/src/denoise.c
@@ -81,6 +81,31 @@ void compute_band_energy(float *bandE, const kiss_fft_cpx *X) {
   }
 }
 
+void compute_band_corr(float *bandE, const kiss_fft_cpx *X, const kiss_fft_cpx *P) {
+  int i;
+  float sum[NB_BANDS] = {0};
+  for (i=0;i<NB_BANDS-1;i++)
+  {
+    int j;
+    int band_size;
+    band_size = (eband5ms[i+1]-eband5ms[i])<<FRAME_SIZE_SHIFT;
+    for (j=0;j<band_size;j++) {
+      float tmp;
+      float frac = (float)j/band_size;
+      tmp = X[(eband5ms[i]<<FRAME_SIZE_SHIFT) + j].r * P[(eband5ms[i]<<FRAME_SIZE_SHIFT) + j].r;
+      tmp += X[(eband5ms[i]<<FRAME_SIZE_SHIFT) + j].i * P[(eband5ms[i]<<FRAME_SIZE_SHIFT) + j].i;
+      sum[i] += (1-frac)*tmp;
+      sum[i+1] += frac*tmp;
+    }
+  }
+  sum[0] *= 2;
+  sum[NB_BANDS-1] *= 2;
+  for (i=0;i<NB_BANDS;i++)
+  {
+    bandE[i] = sum[i];
+  }
+}
+
 void interp_band_gain(float *g, const float *bandE) {
   int i;
   memset(g, 0, FREQ_SIZE);
@@ -225,7 +250,6 @@ DenoiseState *rnnoise_create() {
   return st;
 }
 
-int foo = 0;
 
 static void frame_analysis(DenoiseState *st, kiss_fft_cpx *y, float *Ey, float *features, const float *in) {
   float x[WINDOW_SIZE];
@@ -239,6 +263,7 @@ static void frame_analysis(DenoiseState *st, kiss_fft_cpx *y, float *Ey, float *
   if (1) {
     float p[WINDOW_SIZE];
     kiss_fft_cpx P[WINDOW_SIZE];
+    float Ep[NB_BANDS], Exp[NB_BANDS];
     float pitch_buf[PITCH_BUF_SIZE>>1];
     int pitch_index;
     float gain;
@@ -253,28 +278,20 @@ static void frame_analysis(DenoiseState *st, kiss_fft_cpx *y, float *Ey, float *
 
     gain = remove_doubling(pitch_buf, PITCH_MAX_PERIOD, PITCH_MIN_PERIOD,
             PITCH_FRAME_SIZE, &pitch_index, st->last_period, st->last_gain);
-    if (gain > .9) gain = .9;
-#if 1
-    RNN_MOVE(st->pitch_enh_buf, &st->pitch_enh_buf[FRAME_SIZE], PITCH_BUF_SIZE-FRAME_SIZE);
-    for (i=0;i<FRAME_SIZE;i++)
-      st->pitch_buf[PITCH_BUF_SIZE-FRAME_SIZE+i] = in[i]
-        + ((float)i/FRAME_SIZE)*gain*st->pitch_buf[PITCH_BUF_SIZE-FRAME_SIZE+i-pitch_index]
-        + (1-(float)i/FRAME_SIZE)*st->last_gain*st->pitch_buf[PITCH_BUF_SIZE-FRAME_SIZE+i-st->last_period];
-    RNN_COPY(p, &st->pitch_buf[PITCH_BUF_SIZE-WINDOW_SIZE], WINDOW_SIZE);
-    apply_window(p);
-#else
+    st->last_period = pitch_index;
+    st->last_gain = gain;
     for (i=0;i<WINDOW_SIZE;i++)
       p[i] = st->pitch_buf[PITCH_BUF_SIZE-WINDOW_SIZE-pitch_index+i];
     apply_window(p);
-    for (i=0;i<WINDOW_SIZE;i++)
-      p[i] = .5*(p[i]+x[i]);
-#endif
-    st->last_period = pitch_index;
-    st->last_gain = gain;
-    if (foo) {
-      forward_transform(y, p);
-      compute_band_energy(Ey, y);
+    forward_transform(P, p);
+    compute_band_energy(Ep, P);
+    compute_band_corr(Exp, y, P);
+#if 0
+    if (features) {
+      for (i=0;i<NB_BANDS;i++) printf("%f ", Exp[i]/sqrt(.001+Ey[i]*Ep[i]));
+      printf("\n");
     }
+#endif
   }
   {
     if (features != NULL) {
@@ -377,7 +394,7 @@ int main(int argc, char **argv) {
     fread(tmp, sizeof(short), FRAME_SIZE, f2);
     if (feof(f2)) break;
     for (i=0;i<FRAME_SIZE;i++) n[i] = tmp[i];
-    for (i=0;i<FRAME_SIZE;i++) xn[i] = x[i] + 3*n[i];
+    for (i=0;i<FRAME_SIZE;i++) xn[i] = x[i] + n[i];
     for (i=0;i<FRAME_SIZE;i++) E += x[i]*(float)x[i];
     if (E > 1e9f) {
       vad_cnt=0;
@@ -391,13 +408,12 @@ int main(int argc, char **argv) {
     if (vad_cnt >= 10) vad = 0;
     else if (vad_cnt > 0) vad = 0.5f;
     else vad = 1.f;
-foo=0;
+
     frame_analysis(st, X, Ex, NULL, x);
     frame_analysis(noise_state, N, En, NULL, n);
     for (i=0;i<NB_BANDS;i++) Ln[i] = log10(1e-10+En[i]);
-foo=1;
     frame_analysis(noisy, Y, Ey, features, xn);
-    //printf("%f %d\n", noise_state->last_gain, noise_state->last_period);
+    //printf("%f %d\n", noisy->last_gain, noisy->last_period);
     for (i=0;i<NB_BANDS;i++) {
       g[i] = sqrt((Ex[i]+1e-15)/(Ey[i]+1e-15));
       if (g[i] > 1) g[i] = 1;