Added parameterizable maximum attenuation.

Author: GregorR

Makefile.am

@@ -36,7 +36,7 @@ noinst_PROGRAMS = examples/rnnoise_demo
 endif
 
 examples_rnnoise_demo_SOURCES = examples/rnnoise_demo.c
-examples_rnnoise_demo_LDADD = librnnoise.la
+examples_rnnoise_demo_LDADD = librnnoise.la $(LIBM)
 
 pkgconfigdir = $(libdir)/pkgconfig
 pkgconfig_DATA = rnnoise.pc

examples/rnnoise_demo.c

@@ -24,36 +24,65 @@
    SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */
 
+#include <math.h>
 #include <stdio.h>
 #include "rnnoise.h"
+#include <stdlib.h>
+#include <sys/types.h>
+#include "rnnoise.h"
 
 #define FRAME_SIZE 480
 
 int main(int argc, char **argv) {
-  int i;
+  int i, ci;
   int first = 1;
+  int channels;
   float x[FRAME_SIZE];
-  FILE *f1, *fout;
-  DenoiseState *st;
-  st = rnnoise_create(NULL);
-  if (argc!=3) {
-    fprintf(stderr, "usage: %s <noisy speech> <output denoised>\n", argv[0]);
+  short *tmp;
+  RNNModel *model = NULL;
+  DenoiseState **sts;
+  float max_attenuation;
+  if (argc < 3) {
+    fprintf(stderr, "usage: %s <channels> <max attenuation dB>\n", argv[0]);
+    return 1;
+  }
+
+  channels = atoi(argv[1]);
+  if (channels < 1) channels = 1;
+  max_attenuation = pow(10, -atof(argv[2])/10);
+
+  sts = malloc(channels * sizeof(DenoiseState *));
+  if (!sts) {
+    perror("malloc");
     return 1;
   }
-  f1 = fopen(argv[1], "r");
-  fout = fopen(argv[2], "w");
+  tmp = malloc(channels * FRAME_SIZE * sizeof(short));
+  if (!tmp) {
+      perror("malloc");
+      return 1;
+  }
+  for (i = 0; i < channels; i++) {
+    sts[i] = rnnoise_create(model);
+    rnnoise_set_param(sts[i], RNNOISE_PARAM_MAX_ATTENUATION, max_attenuation);
+  }
+
   while (1) {
-    short tmp[FRAME_SIZE];
-    fread(tmp, sizeof(short), FRAME_SIZE, f1);
-    if (feof(f1)) break;
-    for (i=0;i<FRAME_SIZE;i++) x[i] = tmp[i];
-    rnnoise_process_frame(st, x, x);
-    for (i=0;i<FRAME_SIZE;i++) tmp[i] = x[i];
-    if (!first) fwrite(tmp, sizeof(short), FRAME_SIZE, fout);
+    fread(tmp, sizeof(short), channels * FRAME_SIZE, stdin);
+    if (feof(stdin)) break;
+
+    for (ci = 0; ci < channels; ci++) {
+        for (i=0;i<FRAME_SIZE;i++) x[i] = tmp[i*channels+ci];
+        rnnoise_process_frame(sts[ci], x, x);
+        for (i=0;i<FRAME_SIZE;i++) tmp[i*channels+ci] = x[i];
+    }
+
+    if (!first) fwrite(tmp, sizeof(short), channels * FRAME_SIZE, stdout);
     first = 0;
   }
-  rnnoise_destroy(st);
-  fclose(f1);
-  fclose(fout);
+
+  for (i = 0; i < channels; i++)
+    rnnoise_destroy(sts[i]);
+  free(tmp);
+  free(sts);
   return 0;
 }

include/rnnoise.h

@@ -61,4 +61,9 @@ RNNOISE_EXPORT RNNModel *rnnoise_model_from_file(FILE *f);
 
 RNNOISE_EXPORT void rnnoise_model_free(RNNModel *model);
 
+/* Parameters to a denoise state */
+#define RNNOISE_PARAM_MAX_ATTENUATION           1
+
+RNNOISE_EXPORT void rnnoise_set_param(DenoiseState *st, int param, float value);
+
 #endif

src/denoise.c

@@ -65,6 +65,9 @@
 
 #define NB_FEATURES (NB_BANDS+3*NB_DELTA_CEPS+2)
 
+/* We don't allow max attenuation to be more than 60dB */
+#define MIN_MAX_ATTENUATION 0.000001f
+
 
 #ifndef TRAINING
 #define TRAINING 0
@@ -100,6 +103,8 @@ struct DenoiseState {
   float mem_hp_x[2];
   float lastg[NB_BANDS];
   RNNState rnn;
+
+  float max_attenuation;
 };
 
 #if SMOOTH_BANDS
@@ -508,6 +513,26 @@ float rnnoise_process_frame(DenoiseState *st, float *out, const float *in) {
       g[i] = MAX16(g[i], alpha*st->lastg[i]);
       st->lastg[i] = g[i];
     }
+
+    /* Apply maximum attenuation (minimum value) */
+    if (st->max_attenuation) {
+      float min = 1, mult;
+      for (i=0;i<NB_BANDS;i++) {
+        if (g[i] < min) min = g[i];
+      }
+      if (min < st->max_attenuation) {
+        if (min < MIN_MAX_ATTENUATION)
+          min = MIN_MAX_ATTENUATION;
+        mult = st->max_attenuation / min;
+        for (i=0;i<NB_BANDS;i++) {
+          if (g[i] < MIN_MAX_ATTENUATION) g[i] = MIN_MAX_ATTENUATION;
+          g[i] *= mult;
+          if (g[i] > 1) g[i] = 1;
+          st->lastg[i] = g[i];
+        }
+      }
+    }
+
     interp_band_gain(gf, g);
 #if 1
     for (i=0;i<FREQ_SIZE;i++) {
@@ -521,6 +546,18 @@ float rnnoise_process_frame(DenoiseState *st, float *out, const float *in) {
   return vad_prob;
 }
 
+void rnnoise_set_param(DenoiseState *st, int param, float value)
+{
+  switch (param) {
+    case RNNOISE_PARAM_MAX_ATTENUATION:
+      if ((value > MIN_MAX_ATTENUATION && value <= 1) || value == 0)
+        st->max_attenuation = value;
+      else
+        st->max_attenuation = MIN_MAX_ATTENUATION;
+      break;
+  }
+}
+
 #if TRAINING
 
 static float uni_rand() {