Support model file, training bug fixes and sanity checks

README

@@ -12,6 +12,6 @@ While it is meant to be used as a library, a simple command-line tool is
 provided as an example. It operates on RAW 16-bit (machine endian) mono
 PCM files sampled at 48 kHz. It can be used as:
 
-./examples/rnnoise_demo <number of channels> <maximum attenuation> < input.raw > output.raw
+./examples/rnnoise_demo < input.raw > < output.raw > [< model file >]
 
 The output is also a 16-bit raw PCM file.

TRAINING

@@ -1,3 +1,5 @@
+# NOTE: Required pip packages - 'numpy', 'keras'. Recommended pip packages - 'tensorflow'
+
 (1) cd src ; ./compile.sh
 
 (2) ./denoise_training signal.raw noise.raw count > training.f32
@@ -8,4 +10,4 @@
 
 (4) ./rnn_train.py
 
-(5) ./dump_rnn.py weights.hdf5 ../src/rnn_data.c ../src/rnn_data.h
+(5) ./dump_rnn.py weights.hdf5 ../src/rnn_data.c trained_model.txt orig

examples/rnnoise_demo.c

@@ -34,15 +34,36 @@ int main(int argc, char **argv) {
   int i;
   int first = 1;
   float x[FRAME_SIZE];
+  RNNModel *model = NULL;
   FILE *f1, *fout;
+  FILE *model_fptr;
   DenoiseState *st;
-  st = rnnoise_create(NULL);
-  if (argc!=3) {
-    fprintf(stderr, "usage: %s <noisy speech> <output denoised>\n", argv[0]);
+
+  if (argc < 3) {
+    fprintf(stderr, "usage: %s <noisy speech> <output denoised> [<model file>]\n", argv[0]);
     return 1;
   }
+  if (argc >= 4) {
+    model_fptr = fopen(argv[3], "r");
+    if (!model_fptr) {
+      fprintf(stderr, "Error opening model file \n");
+      return 1;
+    }
+    model = rnnoise_model_from_file(model_fptr);
+    if (!model) {
+      fprintf(stderr, "Model not found \n");
+      return 1;
+    }
+  }
+
+  st = rnnoise_create(model);
   f1 = fopen(argv[1], "r");
   fout = fopen(argv[2], "w");
+  if (!f1) {
+    fprintf(stderr, "Error opening input audio file\n");
+    return 1;
+  }
+
   while (1) {
     short tmp[FRAME_SIZE];
     fread(tmp, sizeof(short), FRAME_SIZE, f1);
@@ -54,7 +75,13 @@ int main(int argc, char **argv) {
     first = 0;
   }
   rnnoise_destroy(st);
+  if (model) {
+    rnnoise_model_free(model);
+  }
   fclose(f1);
   fclose(fout);
+  if (model_fptr) {
+    fclose(model_fptr);
+  }
   return 0;
 }

src/denoise.c

@@ -149,7 +149,7 @@ void compute_band_corr(float *bandE, const kiss_fft_cpx *X, const kiss_fft_cpx *
 
 void interp_band_gain(float *g, const float *bandE) {
   int i;
-  memset(g, 0, FREQ_SIZE);
+  memset(g, 0, FREQ_SIZE * sizeof(float));
   for (i=0;i<NB_BANDS-1;i++)
   {
     int j;
@@ -534,7 +534,15 @@ int main(int argc, char **argv) {
     return 1;
   }
   f1 = fopen(argv[1], "r");
+  if (!f1) {
+      fprintf(stderr, "Error opening speech samples file %s\n", argv[1]);
+      return 1;
+  }
   f2 = fopen(argv[2], "r");
+  if (!f2) {
+      fprintf(stderr, "Error opening noise samples file %s\n", argv[2]);
+      return 1;
+  }
   maxCount = atoi(argv[3]);
   for(i=0;i<150;i++) {
     short tmp[FRAME_SIZE];

training/dump_rnn.py

@@ -37,7 +37,7 @@ def printLayer(f, ft, layer):
     weights = layer.get_weights()
     activation = re.search('function (.*) at', str(layer.activation)).group(1).upper()
     if len(weights) > 2:
-        ft.write('{} {} '.format(weights[0].shape[0], weights[0].shape[1]/3))
+        ft.write('{} {} '.format(weights[0].shape[0], int(weights[0].shape[1]/3)))
     else:
         ft.write('{} {} '.format(weights[0].shape[0], weights[0].shape[1]))
     if activation == 'SIGMOID':
@@ -105,3 +105,4 @@ def mean_squared_sqrt_error(y_true, y_pred):
 #hf.write('};\n')
 
 f.close()
+ft.close()