Reading frames from a raw PCM stream

[pencil]

Hey, I'm using an API that can stream results as either MP3 or raw PCM ("S16LE, 16kHz sample rate"). The following works for reading the results in MP3 format:

container = av.open(result_stream, buffer_size=2048)
audio_stream = None
for stream in container.streams:
    if stream.type == "audio":
        audio_stream = stream
        break

if audio_stream is None:
    raise Exception("No audio stream found")

for frame in container.decode(audio_stream):
    if frame is None:
        break
    // Do something with frame, e.g. resample into desired format

However, when result_stream is in the aforementioned raw PCM format, I get the following error:

av.error.InvalidDataError: [Errno 1094995529] Invalid data found when processing input: '<none>'

I'm guessing I have to somehow tell PyAV the exact format of the input stream since it won't have any container headers indicating sample rate etc. The container_options and/or stream_options sound like what I'm looking for but I can't find any documentation for either of those.

How can I convince PyAV to read frames from a raw PCM stream? Or maybe I just need to read raw bytes myself and construct av.AudioFrames?

[pencil]

I ended up asking ChatGPT (GPT-4) and to my surprise it came up with code that actually works… What a time to be alive.

sample_rate = 16000  # 16kHz
sample_format = "s16"  # S16LE format
channels = 1  # Mono audio
bytes_per_sample = 2  # 16-bit audio
samples_per_frame = 1024  # Number of samples per frame
frame_size = samples_per_frame * channels * bytes_per_sample  # Total frame size in bytes

while True:
    data = result_stream.read(frame_size)
    if not data:
        break

    # Convert byte data to numpy array
    audio_array = np.frombuffer(data, dtype=np.int16).reshape(-1, channels)

    # Change the shape of the array to be (1, n_samples) for a mono packed format
    audio_array = audio_array.reshape(1, -1)

    # Create an AV audio frame
    frame = av.AudioFrame.from_ndarray(audio_array, format=sample_format, layout='mono')
    frame.sample_rate = sample_rate

    # Process the frame
    process_frame(frame)