python-sdks/tests/rtc/test_resampler.py at 7ea29e30b4f3df4e7bf86172daad6351e08a4892 · livekit/python-sdks · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
from livekit.rtc import AudioResampler, AudioResamplerQuality
import time
import wave
import os

# Test fixture directory
FIXTURES_DIR = os.path.join(os.path.dirname(os.path.abspath(__file__)), "fixtures")


def test_audio_resampler():
    wav_file_path = os.path.join(FIXTURES_DIR, "test_audio.wav")

    # Open the wave file
    with wave.open(wav_file_path, "rb") as wf_in:
        n_channels = wf_in.getnchannels()
        sampwidth = wf_in.getsampwidth()
        n_frames = wf_in.getnframes()
        audio_data = bytearray(wf_in.readframes(n_frames))

    if sampwidth != 2:
        raise ValueError(f"Expected 16-bit PCM data, but got {sampwidth * 8}-bit.")

    qualities = [
        AudioResamplerQuality.QUICK,
        AudioResamplerQuality.LOW,
        AudioResamplerQuality.MEDIUM,
        AudioResamplerQuality.HIGH,
        AudioResamplerQuality.VERY_HIGH,
    ]

    for quality in qualities:
        total_time = 0
        nb_runs = 20
        output_frames = []
        for i in range(nb_runs):
            output_frames = []
            resampler = AudioResampler(44100, 8000, quality=quality)

            start_time = time.perf_counter()

            chunk_size = 1024 * n_channels * sampwidth
            output_frames = []
            for i in range(0, len(audio_data), chunk_size):
                chunk = audio_data[i : i + chunk_size]
                frames = resampler.push(bytearray(chunk))
                output_frames.extend(frames)

            frames = resampler.flush()
            output_frames.extend(frames)

            end_time = time.perf_counter()
            total_time += end_time - start_time

        total_time = total_time * 1000
        print(f"Quality: {quality}, Average time: {total_time / nb_runs:.2f}ms")

        output_data = b""

        for frame in output_frames:
            output_data += frame.data

        with wave.open(f"audio_resampled_{quality.name}.wav", "wb") as wf_out:
            wf_out.setnchannels(n_channels)
            wf_out.setsampwidth(sampwidth)
            wf_out.setframerate(8000)
            wf_out.writeframes(output_data)