From 4290c32ae797da4b9da062142d92e63cd646d691 Mon Sep 17 00:00:00 2001
From: Michael Hansen <mike@rhasspy.org>
Date: Fri, 5 Jul 2024 16:13:21 -0500
Subject: [PATCH] Add scripts and main

---
 .gitignore                    |  3 +++
 script/setup                  | 21 ++++++++++++++++
 script/test                   | 13 ++++++++++
 webrtc_noise_gain/__main__.py | 46 +++++++++++++++++++++++++++++++++++
 4 files changed, 83 insertions(+)
 create mode 100755 script/setup
 create mode 100755 script/test
 create mode 100644 webrtc_noise_gain/__main__.py

diff --git a/.gitignore b/.gitignore
index b589e4a..456a380 100644
--- a/.gitignore
+++ b/.gitignore
@@ -14,3 +14,6 @@ __pycache__/
 
 /.venv/
 *.whl
+
+/*.so
+/*.wav
diff --git a/script/setup b/script/setup
new file mode 100755
index 0000000..9269e2a
--- /dev/null
+++ b/script/setup
@@ -0,0 +1,21 @@
+#!/usr/bin/env python3
+import subprocess
+import venv
+from pathlib import Path
+
+_DIR = Path(__file__).parent
+_PROGRAM_DIR = _DIR.parent
+_VENV_DIR = _PROGRAM_DIR / ".venv"
+
+# Create virtual environment
+builder = venv.EnvBuilder(with_pip=True)
+context = builder.ensure_directories(_VENV_DIR)
+builder.create(_VENV_DIR)
+
+# Upgrade dependencies
+pip = [context.env_exe, "-m", "pip"]
+subprocess.check_call(pip + ["install", "--upgrade", "pip"])
+subprocess.check_call(pip + ["install", "--upgrade", "setuptools", "wheel"])
+
+# Install requirements
+subprocess.check_call(pip + ["install", "-r", str(_PROGRAM_DIR / "requirements_dev.txt")])
diff --git a/script/test b/script/test
new file mode 100755
index 0000000..5ad78a8
--- /dev/null
+++ b/script/test
@@ -0,0 +1,13 @@
+#!/usr/bin/env python3
+import subprocess
+import sys
+import venv
+from pathlib import Path
+
+_DIR = Path(__file__).parent
+_PROGRAM_DIR = _DIR.parent
+_VENV_DIR = _PROGRAM_DIR / ".venv"
+_TEST_DIR = _PROGRAM_DIR / "tests"
+
+context = venv.EnvBuilder().ensure_directories(_VENV_DIR)
+subprocess.check_call([context.env_exe, "-m", "pytest", _TEST_DIR] + sys.argv[1:])
diff --git a/webrtc_noise_gain/__main__.py b/webrtc_noise_gain/__main__.py
new file mode 100644
index 0000000..082971c
--- /dev/null
+++ b/webrtc_noise_gain/__main__.py
@@ -0,0 +1,46 @@
+"""Processes a WAV file using the audio processor."""
+
+import argparse
+import wave
+
+from webrtc_noise_gain import AudioProcessor
+
+SAMPLES_10MS = 160
+
+
+def main() -> None:
+    """Main entry point."""
+    parser = argparse.ArgumentParser()
+    parser.add_argument("wav_in", help="WAV input file (16Khz 16-bit mono)")
+    parser.add_argument("wav_out", help="WAV output file")
+    parser.add_argument("--auto-gain", type=int, default=0, choices=range(32))
+    parser.add_argument("--noise-suppression", type=int, default=0, choices=range(5))
+    args = parser.parse_args()
+
+    audio_processor = AudioProcessor(args.auto_gain, args.noise_suppression)
+
+    wav_in: wave.Wave_read = wave.open(args.wav_in, "rb")
+    with wav_in:
+        assert wav_in.getframerate() == 16000, "Sample rate must be 16Khz"
+        assert wav_in.getsampwidth() == 2, "Sample width must be 16-bit"
+        assert wav_in.getnchannels() == 1, "Only mono audio is supported"
+
+        wav_out: wave.Wave_write = wave.open(args.wav_out, "wb")
+        with wav_out:
+            wav_out.setframerate(16000)
+            wav_out.setsampwidth(2)
+            wav_out.setnchannels(1)
+
+            # Process in 10ms chunks
+            chunk = wav_in.readframes(SAMPLES_10MS)
+            while chunk:
+                if len(chunk) < SAMPLES_10MS:
+                    wav_out.writeframes(chunk)
+                    break
+
+                wav_out.writeframes(audio_processor.Process10ms(chunk).audio)
+                chunk = wav_in.readframes(SAMPLES_10MS)
+
+
+if __name__ == "__main__":
+    main()