From 1c0e25a90d12d0359e7b166e907f92e076f49805 Mon Sep 17 00:00:00 2001 From: gui217 Date: Tue, 9 Dec 2025 09:56:20 +0200 Subject: [PATCH] fix unit tests --- tests/test_rnnoise_cancellation.py | 258 +++++++++++++++-------------- tests/test_rnnoise_filter.py | 9 + tests/test_rnnoise_resampling.py | 193 +++++++++++---------- 3 files changed, 232 insertions(+), 228 deletions(-) diff --git a/tests/test_rnnoise_cancellation.py b/tests/test_rnnoise_cancellation.py index 967caa7fa..302dc58e8 100644 --- a/tests/test_rnnoise_cancellation.py +++ b/tests/test_rnnoise_cancellation.py @@ -1,167 +1,169 @@ import asyncio -import os -import wave +import unittest import numpy as np import pytest +try: + import pyrnnoise +except ImportError: + pyrnnoise = None + from pipecat.audio.filters.rnnoise_filter import RNNoiseFilter from pipecat.frames.frames import FilterEnableFrame -async def test_rnnoise_cancellation_functionality(): - print("\nStarting Noise Cancellation Test") +class TestRNNoiseCancellation(unittest.IsolatedAsyncioTestCase): + async def test_rnnoise_cancellation_functionality(self): + print("\nStarting Noise Cancellation Test") - # 1. Check for pyrnnoise - try: - import pyrnnoise - except ImportError: - pytest.skip("pyrnnoise not installed. Cannot verify actual noise cancellation.") - return + # 1. Check for pyrnnoise + if pyrnnoise is None: + self.skipTest("pyrnnoise not installed. Cannot verify actual noise cancellation.") - # 2. Generate clean speech-like audio (Harmonic series) - sample_rate = 48000 - duration = 2.0 - t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False) + # 2. Generate clean speech-like audio (Harmonic series) + sample_rate = 48000 + duration = 2.0 + t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False) - # Fundamental 200Hz + harmonics - clean_signal = np.sin(2 * np.pi * 200 * t) * 0.5 - clean_signal += np.sin(2 * np.pi * 400 * t) * 0.3 - clean_signal += np.sin(2 * np.pi * 600 * t) * 0.2 + # Fundamental 200Hz + harmonics + clean_signal = np.sin(2 * np.pi * 200 * t) * 0.5 + clean_signal += np.sin(2 * np.pi * 400 * t) * 0.3 + clean_signal += np.sin(2 * np.pi * 600 * t) * 0.2 - # Apply envelope to simulate speech (bursts) - # sin(2*pi*2*t) has period 0.5s. - envelope = np.sin(2 * np.pi * 2 * t) - envelope = np.clip(envelope, 0, 1) - clean_signal *= envelope + # Apply envelope to simulate speech (bursts) + # sin(2*pi*2*t) has period 0.5s. + envelope = np.sin(2 * np.pi * 2 * t) + envelope = np.clip(envelope, 0, 1) + clean_signal *= envelope - # 3. Add Noise (White Noise) - noise_level = 0.1 # Reduced noise level slightly to make speech clearer for alignment - noise = np.random.normal(0, noise_level, len(t)) + # 3. Add Noise (White Noise) + noise_level = 0.1 # Reduced noise level slightly to make speech clearer for alignment + noise = np.random.normal(0, noise_level, len(t)) - noisy_signal = clean_signal + noise + noisy_signal = clean_signal + noise - # Normalize to int16 range - noisy_signal = np.clip(noisy_signal, -1, 1) - noisy_int16 = (noisy_signal * 32767).astype(np.int16) - noisy_bytes = noisy_int16.tobytes() + # Normalize to int16 range + noisy_signal = np.clip(noisy_signal, -1, 1) + noisy_int16 = (noisy_signal * 32767).astype(np.int16) + noisy_bytes = noisy_int16.tobytes() - clean_int16 = (clean_signal * 32767).astype(np.int16) + clean_int16 = (clean_signal * 32767).astype(np.int16) - print(f"Generated 2s of noisy audio at {sample_rate}Hz") + print(f"Generated 2s of noisy audio at {sample_rate}Hz") - # 4. Initialize RNNoiseFilter - rnnoise_filter = RNNoiseFilter() - await rnnoise_filter.start(sample_rate) - await rnnoise_filter.process_frame(FilterEnableFrame(enable=True)) + # 4. Initialize RNNoiseFilter + rnnoise_filter = RNNoiseFilter() + await rnnoise_filter.start(sample_rate) + await rnnoise_filter.process_frame(FilterEnableFrame(enable=True)) - # 5. Process - # Feed in chunks - chunk_size = 960 # 20ms - processed_audio = b"" + # 5. Process + # Feed in chunks + chunk_size = 960 # 20ms + processed_audio = b"" - for i in range(0, len(noisy_bytes), chunk_size): - chunk = noisy_bytes[i : i + chunk_size] - result = await rnnoise_filter.filter(chunk) - processed_audio += result + for i in range(0, len(noisy_bytes), chunk_size): + chunk = noisy_bytes[i : i + chunk_size] + result = await rnnoise_filter.filter(chunk) + processed_audio += result - await rnnoise_filter.stop() + await rnnoise_filter.stop() - print(f"Output audio size: {len(processed_audio)}") + print(f"Output audio size: {len(processed_audio)}") - # 6. Verify Noise Reduction - output_int16 = np.frombuffer(processed_audio, dtype=np.int16) + # 6. Verify Noise Reduction + output_int16 = np.frombuffer(processed_audio, dtype=np.int16) - # Truncate to min length - min_len = min(len(clean_int16), len(output_int16)) - clean_trunc = clean_int16[:min_len] - output_trunc = output_int16[:min_len] - noisy_trunc = noisy_int16[:min_len] + # Truncate to min length + min_len = min(len(clean_int16), len(output_int16)) + clean_trunc = clean_int16[:min_len] + output_trunc = output_int16[:min_len] + noisy_trunc = noisy_int16[:min_len] - # 7. Compensate for Delay - # Use cross-correlation on a segment to find delay - # We expect output to be delayed relative to clean (lag is positive) - # search window +/- 2000 samples (~40ms) + # 7. Compensate for Delay + # Use cross-correlation on a segment to find delay + # We expect output to be delayed relative to clean (lag is positive) + # search window +/- 2000 samples (~40ms) - search_range = 2400 # 50ms - # Use the middle of the signal to avoid edge effects and have strong signal - mid_point = min_len // 2 - window_len = 4800 # 100ms + search_range = 2400 # 50ms + # Use the middle of the signal to avoid edge effects and have strong signal + mid_point = min_len // 2 + window_len = 4800 # 100ms - ref_sig = clean_trunc[mid_point : mid_point + window_len].astype(float) - target_sig = output_trunc[ - mid_point - search_range : mid_point + window_len + search_range - ].astype(float) + ref_sig = clean_trunc[mid_point : mid_point + window_len].astype(float) + target_sig = output_trunc[ + mid_point - search_range : mid_point + window_len + search_range + ].astype(float) - correlation = np.correlate(target_sig, ref_sig, mode="valid") - best_idx = np.argmax(correlation) + correlation = np.correlate(target_sig, ref_sig, mode="valid") + best_idx = np.argmax(correlation) - # The 'valid' mode correlation result corresponds to shifts. - # index 0 matches alignment where ref starts at target start. - # target start is (mid_point - search_range). - # ref start is mid_point. - # So index 0 means target is shifted left by search_range (or delay = -search_range). - # delay = best_idx - search_range + # The 'valid' mode correlation result corresponds to shifts. + # index 0 matches alignment where ref starts at target start. + # target start is (mid_point - search_range). + # ref start is mid_point. + # So index 0 means target is shifted left by search_range (or delay = -search_range). + # delay = best_idx - search_range - delay = best_idx - search_range - print(f"Detected delay: {delay} samples ({delay / sample_rate * 1000:.2f} ms)") + delay = best_idx - search_range + print(f"Detected delay: {delay} samples ({delay / sample_rate * 1000:.2f} ms)") - # Shift output to align - if delay > 0: - # Output is delayed, so we need to look at output[delay:] to match clean[0:] - aligned_output = output_trunc[delay:] - aligned_clean = clean_trunc[: len(aligned_output)] - aligned_noisy = noisy_trunc[: len(aligned_output)] - elif delay < 0: - # Output is ahead (unlikely for causal filter), but handling it - aligned_output = output_trunc[:delay] - aligned_clean = clean_trunc[-delay:] - aligned_noisy = noisy_trunc[-delay:] - else: - aligned_output = output_trunc - aligned_clean = clean_trunc - aligned_noisy = noisy_trunc + # Shift output to align + if delay > 0: + # Output is delayed, so we need to look at output[delay:] to match clean[0:] + aligned_output = output_trunc[delay:] + aligned_clean = clean_trunc[: len(aligned_output)] + aligned_noisy = noisy_trunc[: len(aligned_output)] + elif delay < 0: + # Output is ahead (unlikely for causal filter), but handling it + aligned_output = output_trunc[:delay] + aligned_clean = clean_trunc[-delay:] + aligned_noisy = noisy_trunc[-delay:] + else: + aligned_output = output_trunc + aligned_clean = clean_trunc + aligned_noisy = noisy_trunc - # Recalculate MSE on aligned signals - mse_input = np.mean((aligned_noisy.astype(float) - aligned_clean.astype(float)) ** 2) - mse_output = np.mean((aligned_output.astype(float) - aligned_clean.astype(float)) ** 2) + # Recalculate MSE on aligned signals + mse_input = np.mean((aligned_noisy.astype(float) - aligned_clean.astype(float)) ** 2) + mse_output = np.mean((aligned_output.astype(float) - aligned_clean.astype(float)) ** 2) - print(f"MSE (Input vs Clean): {mse_input:.2f}") - print(f"MSE (Output vs Clean): {mse_output:.2f}") + print(f"MSE (Input vs Clean): {mse_input:.2f}") + print(f"MSE (Output vs Clean): {mse_output:.2f}") - # Also check noise reduction in silent regions - # Clean signal envelope is 0 at t=0, 0.25, 0.5... - # Let's find indices where aligned_clean is very small - threshold = 100 # amplitude threshold (out of 32767) - silent_mask = np.abs(aligned_clean) < threshold + # Also check noise reduction in silent regions + # Clean signal envelope is 0 at t=0, 0.25, 0.5... + # Let's find indices where aligned_clean is very small + threshold = 100 # amplitude threshold (out of 32767) + silent_mask = np.abs(aligned_clean) < threshold - if np.sum(silent_mask) > 1000: - noise_power_input = np.mean(aligned_noisy[silent_mask].astype(float) ** 2) - noise_power_output = np.mean(aligned_output[silent_mask].astype(float) ** 2) - print(f"Noise Power in Silence (Input): {noise_power_input:.2f}") - print(f"Noise Power in Silence (Output): {noise_power_output:.2f}") - assert noise_power_output < noise_power_input, "Noise power in silence not reduced" - else: - print("Warning: Not enough silent samples found for noise floor check.") + if np.sum(silent_mask) > 1000: + noise_power_input = np.mean(aligned_noisy[silent_mask].astype(float) ** 2) + noise_power_output = np.mean(aligned_output[silent_mask].astype(float) ** 2) + print(f"Noise Power in Silence (Input): {noise_power_input:.2f}") + print(f"Noise Power in Silence (Output): {noise_power_output:.2f}") + self.assertLess( + noise_power_output, noise_power_input, "Noise power in silence not reduced" + ) + else: + print("Warning: Not enough silent samples found for noise floor check.") - # Main assertion: MSE should improve - # Relax assertion slightly because RNNoise introduces distortion even on clean speech - # But for noisy speech, it should generally be better or at least remove noise. - # If MSE doesn't improve (due to speech distortion), at least Noise Power in Silence should drop. + # Main assertion: MSE should improve + # Relax assertion slightly because RNNoise introduces distortion even on clean speech + # But for noisy speech, it should generally be better or at least remove noise. + # If MSE doesn't improve (due to speech distortion), at least Noise Power in Silence should drop. - if mse_output >= mse_input: - print( - "Warning: Overall MSE did not improve (speech distortion?). Relying on Noise Power check." - ) - # If we passed the noise power check above, we are good. - assert np.sum(silent_mask) > 1000 and np.mean( - aligned_output[silent_mask].astype(float) ** 2 - ) < np.mean(aligned_noisy[silent_mask].astype(float) ** 2) - else: - assert mse_output < mse_input, "MSE did not improve" + if mse_output >= mse_input: + print( + "Warning: Overall MSE did not improve (speech distortion?). Relying on Noise Power check." + ) + # If we passed the noise power check above, we are good. + self.assertTrue( + np.sum(silent_mask) > 1000 + and np.mean(aligned_output[silent_mask].astype(float) ** 2) + < np.mean(aligned_noisy[silent_mask].astype(float) ** 2) + ) + else: + self.assertLess(mse_output, mse_input, "MSE did not improve") - print("Test Passed: Noise cancellation verified.") - - -if __name__ == "__main__": - asyncio.run(test_rnnoise_cancellation_functionality()) + print("Test Passed: Noise cancellation verified.") diff --git a/tests/test_rnnoise_filter.py b/tests/test_rnnoise_filter.py index cfc706af4..6a99b3ba0 100644 --- a/tests/test_rnnoise_filter.py +++ b/tests/test_rnnoise_filter.py @@ -8,11 +8,20 @@ import unittest import numpy as np +try: + import pyrnnoise +except ImportError: + pyrnnoise = None + from pipecat.audio.filters.rnnoise_filter import RNNoiseFilter from pipecat.frames.frames import FilterEnableFrame class TestRNNoiseFilter(unittest.IsolatedAsyncioTestCase): + def setUp(self): + if pyrnnoise is None: + self.skipTest("pyrnnoise not installed") + async def test_rnnoise_filter_reduces_noise(self): """Test that RNNoise filter reduces noise in audio.""" filter = RNNoiseFilter() diff --git a/tests/test_rnnoise_resampling.py b/tests/test_rnnoise_resampling.py index f22413597..acbab4fc0 100644 --- a/tests/test_rnnoise_resampling.py +++ b/tests/test_rnnoise_resampling.py @@ -1,136 +1,129 @@ import asyncio import sys -from unittest.mock import MagicMock +import unittest +from unittest.mock import MagicMock, patch import numpy as np -import pytest -# Mock pyrnnoise BEFORE importing RNNoiseFilter -mock_pyrnnoise = MagicMock() -mock_rnnoise_class = MagicMock() -mock_pyrnnoise.RNNoise = mock_rnnoise_class -sys.modules["pyrnnoise"] = mock_pyrnnoise - -# Now import the filter +# We don't need to mock sys.modules here if we use patch on the imported module member +# But we need to ensure RNNoiseFilter is imported so we can patch its member try: from pipecat.audio.filters.rnnoise_filter import RNNoiseFilter from pipecat.frames.frames import FilterEnableFrame except ImportError as e: + # If dependencies are missing (like numpy?), we can't test print(f"Failed to import RNNoiseFilter: {e}") sys.exit(1) -async def test_rnnoise_resampling_16k_to_48k_and_back(): - print("\nStarting Resampling Test: 16kHz -> 48kHz -> 16kHz") +class TestRNNoiseResampling(unittest.IsolatedAsyncioTestCase): + @patch("pipecat.audio.filters.rnnoise_filter.RNNoise") + async def test_rnnoise_resampling_16k_to_48k_and_back(self, mock_rnnoise_class): + print("\nStarting Resampling Test: 16kHz -> 48kHz -> 16kHz") - # Configure Mock with buffering behavior - processed_chunks_count = 0 - buffer = np.array([], dtype=np.int16) + # Configure Mock with buffering behavior + processed_chunks_count = 0 + buffer = np.array([], dtype=np.int16) - def side_effect_process_chunk(audio_samples, partial=False): - nonlocal buffer, processed_chunks_count + def side_effect_process_chunk(audio_samples, partial=False): + nonlocal buffer, processed_chunks_count - # Append new samples to buffer - if len(audio_samples) > 0: - buffer = np.concatenate((buffer, audio_samples)) + # Append new samples to buffer + if len(audio_samples) > 0: + buffer = np.concatenate((buffer, audio_samples)) - # Yield 480-sample chunks - while len(buffer) >= 480: - chunk = buffer[:480] - buffer = buffer[480:] - processed_chunks_count += 1 + # Yield 480-sample chunks + while len(buffer) >= 480: + chunk = buffer[:480] + buffer = buffer[480:] + processed_chunks_count += 1 - # Simulate processing (pass through) - # Convert int16 -> float32 [-1, 1] - normalized = chunk.astype(np.float32) / 32768.0 - yield 0.99, normalized + # Simulate processing (pass through) + # Convert int16 -> float32 [-1, 1] + normalized = chunk.astype(np.float32) / 32768.0 + yield 0.99, normalized - mock_rnnoise_instance = MagicMock() - mock_rnnoise_instance.denoise_chunk.side_effect = side_effect_process_chunk - mock_rnnoise_class.return_value = mock_rnnoise_instance + mock_rnnoise_instance = MagicMock() + mock_rnnoise_instance.denoise_chunk.side_effect = side_effect_process_chunk + mock_rnnoise_class.return_value = mock_rnnoise_instance - # 1. Generate 1 second of 16kHz audio (sine wave 440Hz) - sample_rate = 16000 - duration = 1.0 - t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False) - audio_data = (np.sin(2 * np.pi * 440 * t) * 32767).astype(np.int16) - audio_bytes = audio_data.tobytes() + # 1. Generate 1 second of 16kHz audio (sine wave 440Hz) + sample_rate = 16000 + duration = 1.0 + t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False) + audio_data = (np.sin(2 * np.pi * 440 * t) * 32767).astype(np.int16) + audio_bytes = audio_data.tobytes() - print(f"Input audio: {len(audio_bytes)} bytes, {len(audio_data)} samples at {sample_rate}Hz") + print( + f"Input audio: {len(audio_bytes)} bytes, {len(audio_data)} samples at {sample_rate}Hz" + ) - # 2. Initialize RNNoiseFilter - rnnoise_filter = RNNoiseFilter() - await rnnoise_filter.start(sample_rate) + # 2. Initialize RNNoiseFilter + # This will use the patched RNNoise + rnnoise_filter = RNNoiseFilter() + await rnnoise_filter.start(sample_rate) - # Enable filtering - await rnnoise_filter.process_frame(FilterEnableFrame(enable=True)) + # Enable filtering + await rnnoise_filter.process_frame(FilterEnableFrame(enable=True)) - # 3. Process audio in chunks - chunk_size = 320 # 160 samples (10ms at 16k) * 2 bytes - processed_audio = b"" + # 3. Process audio in chunks + chunk_size = 320 # 160 samples (10ms at 16k) * 2 bytes + processed_audio = b"" - for i in range(0, len(audio_bytes), chunk_size): - chunk = audio_bytes[i : i + chunk_size] - result = await rnnoise_filter.filter(chunk) - processed_audio += result + for i in range(0, len(audio_bytes), chunk_size): + chunk = audio_bytes[i : i + chunk_size] + result = await rnnoise_filter.filter(chunk) + processed_audio += result - await rnnoise_filter.stop() + await rnnoise_filter.stop() - print(f"Output audio: {len(processed_audio)} bytes") - print(f"Processed chunks (internal 480 samples): {processed_chunks_count}") + print(f"Output audio: {len(processed_audio)} bytes") + print(f"Processed chunks (internal 480 samples): {processed_chunks_count}") - # 4. Verify output length - # Expect roughly same length - # Input: 16000 samples. - # Upsampled to 48000. - # 48000 / 480 = 100 chunks. - # So we expect roughly 100 calls to process_chunk. - expected_chunks = (len(audio_data) * 48000 // sample_rate) // 480 - print(f"Expected chunks: ~{expected_chunks}") + # 4. Verify output length + # Expect roughly same length + expected_chunks = (len(audio_data) * 48000 // sample_rate) // 480 + print(f"Expected chunks: ~{expected_chunks}") - # Check that we actually processed something - assert processed_chunks_count >= expected_chunks - 5, "Too few chunks processed" + # Check that we actually processed something + self.assertGreaterEqual( + processed_chunks_count, expected_chunks - 5, "Too few chunks processed" + ) - # Check output length - assert len(processed_audio) > 0, "Output should not be empty" + # Check output length + self.assertGreater(len(processed_audio), 0, "Output should not be empty") - # Check length matches input (with some tolerance for buffering latency) - # Since we don't flush the filter explicitly (no flush method in RNNoiseFilter yet), - # some data might remain in buffers. - # Max loss: - # - Resampler input buffer - # - RNNoise buffer (max 480 samples = 10ms) - # - Resampler output buffer + # Check length matches input (with some tolerance for buffering latency) + # 100ms tolerance? + byte_tolerance = int(0.2 * sample_rate * 2) + self.assertGreaterEqual( + len(processed_audio), + len(audio_bytes) - byte_tolerance, + f"Output too short: {len(processed_audio)} vs {len(audio_bytes)}", + ) + self.assertLessEqual( + len(processed_audio), + len(audio_bytes) + byte_tolerance, + f"Output too long: {len(processed_audio)} vs {len(audio_bytes)}", + ) - # 100ms tolerance? - byte_tolerance = int(0.2 * sample_rate * 2) - assert len(processed_audio) >= len(audio_bytes) - byte_tolerance, ( - f"Output too short: {len(processed_audio)} vs {len(audio_bytes)}" - ) - assert len(processed_audio) <= len(audio_bytes) + byte_tolerance, ( - f"Output too long: {len(processed_audio)} vs {len(audio_bytes)}" - ) + # 5. Check sample rate / pitch preservation + output_data = np.frombuffer(processed_audio, dtype=np.int16) - # 5. Check sample rate / pitch preservation - # If we upsampled and downsampled correctly, the pitch should be 440Hz. - output_data = np.frombuffer(processed_audio, dtype=np.int16) + if len(output_data) > 2000: + # Use a window in the middle + start_idx = len(output_data) // 4 + end_idx = 3 * len(output_data) // 4 + segment = output_data[start_idx:end_idx] - if len(output_data) > 2000: - # Use a window in the middle - start_idx = len(output_data) // 4 - end_idx = 3 * len(output_data) // 4 - segment = output_data[start_idx:end_idx] + fft = np.fft.rfft(segment) + freqs = np.fft.rfftfreq(len(segment), d=1 / sample_rate) + peak_idx = np.argmax(np.abs(fft)) + peak_freq = freqs[peak_idx] - fft = np.fft.rfft(segment) - freqs = np.fft.rfftfreq(len(segment), d=1 / sample_rate) - peak_idx = np.argmax(np.abs(fft)) - peak_freq = freqs[peak_idx] + print(f"Peak frequency: {peak_freq:.2f} Hz") + self.assertLess( + abs(peak_freq - 440), 50, f"Frequency shifted significantly: {peak_freq} vs 440" + ) - print(f"Peak frequency: {peak_freq:.2f} Hz") - assert abs(peak_freq - 440) < 50, f"Frequency shifted significantly: {peak_freq} vs 440" - - print("Test Passed: Resampling logic verified (with mocked RNNoise).") - - -if __name__ == "__main__": - asyncio.run(test_rnnoise_resampling_16k_to_48k_and_back()) + print("Test Passed: Resampling logic verified (with mocked RNNoise).")