wx44wx/pipecat
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix unit tests

Browse Source
This commit is contained in:
gui217
2025-12-09 09:56:20 +02:00
parent c48858742a
commit 1c0e25a90d
3 changed files with 232 additions and 228 deletions
Download Patch File Download Diff File Expand all files Collapse all files

258
tests/test_rnnoise_cancellation.py
View File

@@ -1,167 +1,169 @@
import asyncio import asyncio
import os import unittest
import wave
import numpy as np import numpy as np
import pytest import pytest
try:
import pyrnnoise
except ImportError:
pyrnnoise = None
from pipecat.audio.filters.rnnoise_filter import RNNoiseFilter from pipecat.audio.filters.rnnoise_filter import RNNoiseFilter
from pipecat.frames.frames import FilterEnableFrame from pipecat.frames.frames import FilterEnableFrame
async def test_rnnoise_cancellation_functionality(): class TestRNNoiseCancellation(unittest.IsolatedAsyncioTestCase):
print("\nStarting Noise Cancellation Test") async def test_rnnoise_cancellation_functionality(self):
print("\nStarting Noise Cancellation Test")
# 1. Check for pyrnnoise # 1. Check for pyrnnoise
try: if pyrnnoise is None:
import pyrnnoise self.skipTest("pyrnnoise not installed. Cannot verify actual noise cancellation.")
except ImportError:
pytest.skip("pyrnnoise not installed. Cannot verify actual noise cancellation.")
return
# 2. Generate clean speech-like audio (Harmonic series) # 2. Generate clean speech-like audio (Harmonic series)
sample_rate = 48000 sample_rate = 48000
duration = 2.0 duration = 2.0
t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False) t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False)
# Fundamental 200Hz + harmonics # Fundamental 200Hz + harmonics
clean_signal = np.sin(2 * np.pi * 200 * t) * 0.5 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 * 400 * t) * 0.3
clean_signal += np.sin(2 * np.pi * 600 * t) * 0.2 clean_signal += np.sin(2 * np.pi * 600 * t) * 0.2
# Apply envelope to simulate speech (bursts) # Apply envelope to simulate speech (bursts)
# sin(2*pi*2*t) has period 0.5s. # sin(2*pi*2*t) has period 0.5s.
envelope = np.sin(2 * np.pi * 2 * t) envelope = np.sin(2 * np.pi * 2 * t)
envelope = np.clip(envelope, 0, 1) envelope = np.clip(envelope, 0, 1)
clean_signal *= envelope clean_signal *= envelope
# 3. Add Noise (White Noise) # 3. Add Noise (White Noise)
noise_level = 0.1 # Reduced noise level slightly to make speech clearer for alignment noise_level = 0.1 # Reduced noise level slightly to make speech clearer for alignment
noise = np.random.normal(0, noise_level, len(t)) noise = np.random.normal(0, noise_level, len(t))
noisy_signal = clean_signal + noise noisy_signal = clean_signal + noise
# Normalize to int16 range # Normalize to int16 range
noisy_signal = np.clip(noisy_signal, -1, 1) noisy_signal = np.clip(noisy_signal, -1, 1)
noisy_int16 = (noisy_signal * 32767).astype(np.int16) noisy_int16 = (noisy_signal * 32767).astype(np.int16)
noisy_bytes = noisy_int16.tobytes() 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 # 4. Initialize RNNoiseFilter
rnnoise_filter = RNNoiseFilter() rnnoise_filter = RNNoiseFilter()
await rnnoise_filter.start(sample_rate) await rnnoise_filter.start(sample_rate)
await rnnoise_filter.process_frame(FilterEnableFrame(enable=True)) await rnnoise_filter.process_frame(FilterEnableFrame(enable=True))
# 5. Process # 5. Process
# Feed in chunks # Feed in chunks
chunk_size = 960 # 20ms chunk_size = 960 # 20ms
processed_audio = b"" processed_audio = b""
for i in range(0, len(noisy_bytes), chunk_size): for i in range(0, len(noisy_bytes), chunk_size):
chunk = noisy_bytes[i : i + chunk_size] chunk = noisy_bytes[i : i + chunk_size]
result = await rnnoise_filter.filter(chunk) result = await rnnoise_filter.filter(chunk)
processed_audio += result 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 # 6. Verify Noise Reduction
output_int16 = np.frombuffer(processed_audio, dtype=np.int16) output_int16 = np.frombuffer(processed_audio, dtype=np.int16)
# Truncate to min length # Truncate to min length
min_len = min(len(clean_int16), len(output_int16)) min_len = min(len(clean_int16), len(output_int16))
clean_trunc = clean_int16[:min_len] clean_trunc = clean_int16[:min_len]
output_trunc = output_int16[:min_len] output_trunc = output_int16[:min_len]
noisy_trunc = noisy_int16[:min_len] noisy_trunc = noisy_int16[:min_len]
# 7. Compensate for Delay # 7. Compensate for Delay
# Use cross-correlation on a segment to find delay # Use cross-correlation on a segment to find delay
# We expect output to be delayed relative to clean (lag is positive) # We expect output to be delayed relative to clean (lag is positive)
# search window +/- 2000 samples (~40ms) # search window +/- 2000 samples (~40ms)
search_range = 2400 # 50ms search_range = 2400 # 50ms
# Use the middle of the signal to avoid edge effects and have strong signal # Use the middle of the signal to avoid edge effects and have strong signal
mid_point = min_len // 2 mid_point = min_len // 2
window_len = 4800 # 100ms window_len = 4800 # 100ms
ref_sig = clean_trunc[mid_point : mid_point + window_len].astype(float) ref_sig = clean_trunc[mid_point : mid_point + window_len].astype(float)
target_sig = output_trunc[ target_sig = output_trunc[
mid_point - search_range : mid_point + window_len + search_range mid_point - search_range : mid_point + window_len + search_range
].astype(float) ].astype(float)
correlation = np.correlate(target_sig, ref_sig, mode="valid") correlation = np.correlate(target_sig, ref_sig, mode="valid")
best_idx = np.argmax(correlation) best_idx = np.argmax(correlation)
# The 'valid' mode correlation result corresponds to shifts. # The 'valid' mode correlation result corresponds to shifts.
# index 0 matches alignment where ref starts at target start. # index 0 matches alignment where ref starts at target start.
# target start is (mid_point - search_range). # target start is (mid_point - search_range).
# ref start is mid_point. # ref start is mid_point.
# So index 0 means target is shifted left by search_range (or delay = -search_range). # 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
delay = best_idx - search_range delay = best_idx - search_range
print(f"Detected delay: {delay} samples ({delay / sample_rate * 1000:.2f} ms)") print(f"Detected delay: {delay} samples ({delay / sample_rate * 1000:.2f} ms)")
# Shift output to align # Shift output to align
if delay > 0: if delay > 0:
# Output is delayed, so we need to look at output[delay:] to match clean[0:] # Output is delayed, so we need to look at output[delay:] to match clean[0:]
aligned_output = output_trunc[delay:] aligned_output = output_trunc[delay:]
aligned_clean = clean_trunc[: len(aligned_output)] aligned_clean = clean_trunc[: len(aligned_output)]
aligned_noisy = noisy_trunc[: len(aligned_output)] aligned_noisy = noisy_trunc[: len(aligned_output)]
elif delay < 0: elif delay < 0:
# Output is ahead (unlikely for causal filter), but handling it # Output is ahead (unlikely for causal filter), but handling it
aligned_output = output_trunc[:delay] aligned_output = output_trunc[:delay]
aligned_clean = clean_trunc[-delay:] aligned_clean = clean_trunc[-delay:]
aligned_noisy = noisy_trunc[-delay:] aligned_noisy = noisy_trunc[-delay:]
else: else:
aligned_output = output_trunc aligned_output = output_trunc
aligned_clean = clean_trunc aligned_clean = clean_trunc
aligned_noisy = noisy_trunc aligned_noisy = noisy_trunc
# Recalculate MSE on aligned signals # Recalculate MSE on aligned signals
mse_input = np.mean((aligned_noisy.astype(float) - aligned_clean.astype(float)) ** 2) 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) 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 (Input vs Clean): {mse_input:.2f}")
print(f"MSE (Output vs Clean): {mse_output:.2f}") print(f"MSE (Output vs Clean): {mse_output:.2f}")
# Also check noise reduction in silent regions # Also check noise reduction in silent regions
# Clean signal envelope is 0 at t=0, 0.25, 0.5... # Clean signal envelope is 0 at t=0, 0.25, 0.5...
# Let's find indices where aligned_clean is very small # Let's find indices where aligned_clean is very small
threshold = 100 # amplitude threshold (out of 32767) threshold = 100 # amplitude threshold (out of 32767)
silent_mask = np.abs(aligned_clean) < threshold silent_mask = np.abs(aligned_clean) < threshold
if np.sum(silent_mask) > 1000: if np.sum(silent_mask) > 1000:
noise_power_input = np.mean(aligned_noisy[silent_mask].astype(float) ** 2) noise_power_input = np.mean(aligned_noisy[silent_mask].astype(float) ** 2)
noise_power_output = np.mean(aligned_output[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 (Input): {noise_power_input:.2f}")
print(f"Noise Power in Silence (Output): {noise_power_output:.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" self.assertLess(
else: noise_power_output, noise_power_input, "Noise power in silence not reduced"
print("Warning: Not enough silent samples found for noise floor check.") )
else:
print("Warning: Not enough silent samples found for noise floor check.")
# Main assertion: MSE should improve # Main assertion: MSE should improve
# Relax assertion slightly because RNNoise introduces distortion even on clean speech # 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. # 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 doesn't improve (due to speech distortion), at least Noise Power in Silence should drop.
if mse_output >= mse_input: if mse_output >= mse_input:
print( print(
"Warning: Overall MSE did not improve (speech distortion?). Relying on Noise Power check." "Warning: Overall MSE did not improve (speech distortion?). Relying on Noise Power check."
) )
# If we passed the noise power check above, we are good. # If we passed the noise power check above, we are good.
assert np.sum(silent_mask) > 1000 and np.mean( self.assertTrue(
aligned_output[silent_mask].astype(float) ** 2 np.sum(silent_mask) > 1000
) < np.mean(aligned_noisy[silent_mask].astype(float) ** 2) and np.mean(aligned_output[silent_mask].astype(float) ** 2)
else: < np.mean(aligned_noisy[silent_mask].astype(float) ** 2)
assert mse_output < mse_input, "MSE did not improve" )
else:
self.assertLess(mse_output, mse_input, "MSE did not improve")
print("Test Passed: Noise cancellation verified.") print("Test Passed: Noise cancellation verified.")
if __name__ == "__main__":
asyncio.run(test_rnnoise_cancellation_functionality())

9
tests/test_rnnoise_filter.py
View File

@@ -8,11 +8,20 @@ import unittest
import numpy as np import numpy as np
try:
import pyrnnoise
except ImportError:
pyrnnoise = None
from pipecat.audio.filters.rnnoise_filter import RNNoiseFilter from pipecat.audio.filters.rnnoise_filter import RNNoiseFilter
from pipecat.frames.frames import FilterEnableFrame from pipecat.frames.frames import FilterEnableFrame
class TestRNNoiseFilter(unittest.IsolatedAsyncioTestCase): class TestRNNoiseFilter(unittest.IsolatedAsyncioTestCase):
def setUp(self):
if pyrnnoise is None:
self.skipTest("pyrnnoise not installed")
async def test_rnnoise_filter_reduces_noise(self): async def test_rnnoise_filter_reduces_noise(self):
"""Test that RNNoise filter reduces noise in audio.""" """Test that RNNoise filter reduces noise in audio."""
filter = RNNoiseFilter() filter = RNNoiseFilter()

193
tests/test_rnnoise_resampling.py
View File

@@ -1,136 +1,129 @@
import asyncio import asyncio
import sys import sys
from unittest.mock import MagicMock import unittest
from unittest.mock import MagicMock, patch
import numpy as np import numpy as np
import pytest
# Mock pyrnnoise BEFORE importing RNNoiseFilter # We don't need to mock sys.modules here if we use patch on the imported module member
mock_pyrnnoise = MagicMock() # But we need to ensure RNNoiseFilter is imported so we can patch its member
mock_rnnoise_class = MagicMock()
mock_pyrnnoise.RNNoise = mock_rnnoise_class
sys.modules["pyrnnoise"] = mock_pyrnnoise
# Now import the filter
try: try:
from pipecat.audio.filters.rnnoise_filter import RNNoiseFilter from pipecat.audio.filters.rnnoise_filter import RNNoiseFilter
from pipecat.frames.frames import FilterEnableFrame from pipecat.frames.frames import FilterEnableFrame
except ImportError as e: except ImportError as e:
# If dependencies are missing (like numpy?), we can't test
print(f"Failed to import RNNoiseFilter: {e}") print(f"Failed to import RNNoiseFilter: {e}")
sys.exit(1) sys.exit(1)
async def test_rnnoise_resampling_16k_to_48k_and_back(): class TestRNNoiseResampling(unittest.IsolatedAsyncioTestCase):
print("\nStarting Resampling Test: 16kHz -> 48kHz -> 16kHz") @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 # Configure Mock with buffering behavior
processed_chunks_count = 0 processed_chunks_count = 0
buffer = np.array([], dtype=np.int16) buffer = np.array([], dtype=np.int16)
def side_effect_process_chunk(audio_samples, partial=False): def side_effect_process_chunk(audio_samples, partial=False):
nonlocal buffer, processed_chunks_count nonlocal buffer, processed_chunks_count
# Append new samples to buffer # Append new samples to buffer
if len(audio_samples) > 0: if len(audio_samples) > 0:
buffer = np.concatenate((buffer, audio_samples)) buffer = np.concatenate((buffer, audio_samples))
# Yield 480-sample chunks # Yield 480-sample chunks
while len(buffer) >= 480: while len(buffer) >= 480:
chunk = buffer[:480] chunk = buffer[:480]
buffer = buffer[480:] buffer = buffer[480:]
processed_chunks_count += 1 processed_chunks_count += 1
# Simulate processing (pass through) # Simulate processing (pass through)
# Convert int16 -> float32 [-1, 1] # Convert int16 -> float32 [-1, 1]
normalized = chunk.astype(np.float32) / 32768.0 normalized = chunk.astype(np.float32) / 32768.0
yield 0.99, normalized yield 0.99, normalized
mock_rnnoise_instance = MagicMock() mock_rnnoise_instance = MagicMock()
mock_rnnoise_instance.denoise_chunk.side_effect = side_effect_process_chunk mock_rnnoise_instance.denoise_chunk.side_effect = side_effect_process_chunk
mock_rnnoise_class.return_value = mock_rnnoise_instance mock_rnnoise_class.return_value = mock_rnnoise_instance
# 1. Generate 1 second of 16kHz audio (sine wave 440Hz) # 1. Generate 1 second of 16kHz audio (sine wave 440Hz)
sample_rate = 16000 sample_rate = 16000
duration = 1.0 duration = 1.0
t = np.linspace(0, duration, int(sample_rate * duration), endpoint=False) 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_data = (np.sin(2 * np.pi * 440 * t) * 32767).astype(np.int16)
audio_bytes = audio_data.tobytes() 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 # 2. Initialize RNNoiseFilter
rnnoise_filter = RNNoiseFilter() # This will use the patched RNNoise
await rnnoise_filter.start(sample_rate) rnnoise_filter = RNNoiseFilter()
await rnnoise_filter.start(sample_rate)
# Enable filtering # Enable filtering
await rnnoise_filter.process_frame(FilterEnableFrame(enable=True)) await rnnoise_filter.process_frame(FilterEnableFrame(enable=True))
# 3. Process audio in chunks # 3. Process audio in chunks
chunk_size = 320 # 160 samples (10ms at 16k) * 2 bytes chunk_size = 320 # 160 samples (10ms at 16k) * 2 bytes
processed_audio = b"" processed_audio = b""
for i in range(0, len(audio_bytes), chunk_size): for i in range(0, len(audio_bytes), chunk_size):
chunk = audio_bytes[i : i + chunk_size] chunk = audio_bytes[i : i + chunk_size]
result = await rnnoise_filter.filter(chunk) result = await rnnoise_filter.filter(chunk)
processed_audio += result processed_audio += result
await rnnoise_filter.stop() await rnnoise_filter.stop()
print(f"Output audio: {len(processed_audio)} bytes") print(f"Output audio: {len(processed_audio)} bytes")
print(f"Processed chunks (internal 480 samples): {processed_chunks_count}") print(f"Processed chunks (internal 480 samples): {processed_chunks_count}")
# 4. Verify output length # 4. Verify output length
# Expect roughly same length # Expect roughly same length
# Input: 16000 samples. expected_chunks = (len(audio_data) * 48000 // sample_rate) // 480
# Upsampled to 48000. print(f"Expected chunks: ~{expected_chunks}")
# 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}")
# Check that we actually processed something # Check that we actually processed something
assert processed_chunks_count >= expected_chunks - 5, "Too few chunks processed" self.assertGreaterEqual(
processed_chunks_count, expected_chunks - 5, "Too few chunks processed"
)
# Check output length # Check output length
assert len(processed_audio) > 0, "Output should not be empty" self.assertGreater(len(processed_audio), 0, "Output should not be empty")
# Check length matches input (with some tolerance for buffering latency) # Check length matches input (with some tolerance for buffering latency)
# Since we don't flush the filter explicitly (no flush method in RNNoiseFilter yet), # 100ms tolerance?
# some data might remain in buffers. byte_tolerance = int(0.2 * sample_rate * 2)
# Max loss: self.assertGreaterEqual(
# - Resampler input buffer len(processed_audio),
# - RNNoise buffer (max 480 samples = 10ms) len(audio_bytes) - byte_tolerance,
# - Resampler output buffer 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? # 5. Check sample rate / pitch preservation
byte_tolerance = int(0.2 * sample_rate * 2) output_data = np.frombuffer(processed_audio, dtype=np.int16)
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 if len(output_data) > 2000:
# If we upsampled and downsampled correctly, the pitch should be 440Hz. # Use a window in the middle
output_data = np.frombuffer(processed_audio, dtype=np.int16) 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: fft = np.fft.rfft(segment)
# Use a window in the middle freqs = np.fft.rfftfreq(len(segment), d=1 / sample_rate)
start_idx = len(output_data) // 4 peak_idx = np.argmax(np.abs(fft))
end_idx = 3 * len(output_data) // 4 peak_freq = freqs[peak_idx]
segment = output_data[start_idx:end_idx]
fft = np.fft.rfft(segment) print(f"Peak frequency: {peak_freq:.2f} Hz")
freqs = np.fft.rfftfreq(len(segment), d=1 / sample_rate) self.assertLess(
peak_idx = np.argmax(np.abs(fft)) abs(peak_freq - 440), 50, f"Frequency shifted significantly: {peak_freq} vs 440"
peak_freq = freqs[peak_idx] )
print(f"Peak frequency: {peak_freq:.2f} Hz") print("Test Passed: Resampling logic verified (with mocked RNNoise).")
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())