inital commit

1 files changed, 59 insertions, 0 deletions

diff --git a/lib/tests/test_engines.py b/lib/tests/test_engines.py
new file mode 100644
index 0000000..0efccc6
--- /dev/null
+++ b/lib/tests/test_engines.py
@@ -0,0 +1,59 @@
+import os
+
+import numpy as np
+import pytest
+
+
+def _metrics_floor(x, sr):
+    from producer import meters
+
+    return meters.noise_floor_db(x, sr)
+
+
+@pytest.mark.slow
+def test_dfn3_reduces_noise(sr, noisy_speech):
+    pytest.importorskip("torch")
+    pytest.importorskip("df")
+    from producer.engines import denoise_dfn
+
+    x = noisy_speech[: sr * 4]
+    y, eng, _dev = denoise_dfn.denoise(x, sr, 1.0, "cpu")
+    assert "dfn" in eng
+    assert _metrics_floor(y, sr) < _metrics_floor(x, sr) - 5.0
+    assert np.corrcoef(x, y.astype(np.float64))[0, 1] > 0.9
+
+
+@pytest.mark.slow
+def test_zipenhancer_reduces_noise(sr, noisy_speech):
+    pytest.importorskip("torch")
+    pytest.importorskip("zipenhancer")
+    from producer.engines import denoise_zip
+
+    x = noisy_speech[: sr * 4]
+    y, eng, _ = denoise_zip.denoise(x, sr, 1.0, "cpu")
+    assert "zipenhancer" in eng
+    assert _metrics_floor(y, sr) < _metrics_floor(x, sr) - 5.0
+
+
+@pytest.mark.slow
+def test_mossformer2_enhances(sr, noisy_speech):
+    pytest.importorskip("torch")
+    pytest.importorskip("clearvoice")
+    from producer.engines import enhance_mossformer
+
+    x = noisy_speech[: sr * 4]
+    y, eng, _ = enhance_mossformer.enhance(x, sr, 1.0, "cpu")
+    assert "mossformer2" in eng
+    assert np.all(np.isfinite(y))
+
+
+@pytest.mark.slow
+def test_resemble_enhance(sr, noisy_speech):
+    if not os.environ.get("PRODUCER_TEST_RESEMBLE"):
+        pytest.skip("set PRODUCER_TEST_RESEMBLE=1 to run the isolated-venv generative engine")
+    from producer.engines import enhance_resemble
+
+    x = noisy_speech[: sr * 4]
+    y, eng, _ = enhance_resemble.enhance(x, sr, 1.0, "cpu")
+    assert "resemble" in eng
+    assert np.all(np.isfinite(y))