How to combine streamlit_webrtc and asyncio for parrallel audio capture and inference?

[SuperKogito]

How can I combine streamlit_webrtc and asyncio to run inference in parallel with the audio buffering? is there any example of how to do this?

I am currently trying to deploy some Essentia model for music genre recognition with the help of streamlit_webrtc. Basically I want to capture the audio stream and then when a certain amount of frames is reached (e.g. 5 seconds) I run the prediction. Unfortunately is a lot of lags and delays so I am trying to run these in parallel with the help of asyncio. Unfortunately, this seems to be blocking my pipeline, possibly because streamlit_webrtc uses asyncio too. So what is the proper way to do this. Here is my code:

import json
import time
import json
import pydub
import queue
import socket
import joblib
import librosa
import asyncio
import logging
import numpy as np
import pandas as pd
import logging.handlers
import streamlit as st
import tensorflow as tf
from pathlib import Path
from pprint import pprint
from streamlit_webrtc import WebRtcMode, webrtc_streamer
from essentia.standard import TensorflowPredictEffnetDiscogs, TensorflowPredict2D

DEBUG = False
SAMPLE_RATE = 16000
HERE = Path(__file__).parent
logger = logging.getLogger(__name__)

# Disable eager mode for tf.v1 compatibility with tf.v2
tf.compat.v1.disable_eager_execution()

import warnings
warnings.filterwarnings('ignore')

def get_metadata(fname, section=None):
    with open(fname) as f:
        meta = json.load(f)

    if section:
        return meta[section]
    return meta

MODEL_LABELS = {"ESSENTIA_genre_discogs400-discogs-effnet-1-random-forrest": get_metadata("pretrained-models/genre_discogs400-discogs-effnet-1.json", "classes")}

def load_essentia_model(model, labels, n_frames=None, overlap=None, debug=True):
    embeddings_model_args = {"graphFilename": "pretrained-models/discogs-effnet-bs64-1.pb", "output": "PartitionedCall:1"}
    predictions_model_args = {"graphFilename": "pretrained-models/genre_discogs400-discogs-effnet-1.pb", "input": "serving_default_model_Placeholder", "output": "PartitionedCall:0"}
    embeddings_model = TensorflowPredictEffnetDiscogs(**embeddings_model_args)
    predictions_model = TensorflowPredict2D(**predictions_model_args)
    return embeddings_model, predictions_model

def predict_essentia(audio, embeddings_model, predictions_model):
    # Run the models
    #print("-> Audio to predict: ", audio.shape, audio.dtype, type(audio))
    embeddings = embeddings_model(audio)
    predictions = predictions_model(embeddings)
    return predictions       

async def audio_collector(webrtc_ctx, sound_chunk_queue):
    """Asynchronously collect audio frames from webrtc context"""
    while webrtc_ctx.state.playing:
        try:
            audio_frames = webrtc_ctx.audio_receiver.get_frames(timeout=1)
        except queue.Empty:
            continue
        
        for audio_frame in audio_frames:
            sound = pydub.AudioSegment(
                data=audio_frame.to_ndarray().tobytes(),
                sample_width=audio_frame.format.bytes,
                frame_rate=audio_frame.sample_rate,
                channels=len(audio_frame.layout.channels)
            )
            # Put the sound into the queue
            await sound_chunk_queue.put(sound)

async def run_inference(sound_chunk, model, embeddings_model, predictions_model, random_forrest_model, text_output, json_output):
    start_process = time.time()

    # Audio processing
    print(f"Inference: Received sound_chunk of length {len(sound_chunk)}")
    sound_chunk = sound_chunk.set_channels(1).set_frame_rate(SAMPLE_RATE)
    musicnn_audio_buffer = np.array(sound_chunk.get_array_of_samples(), dtype=np.int16)
    float_musicnn_audio_buffer = musicnn_audio_buffer.astype(np.float32, order='C') / 32768.0

    # Run inference
    likelihoods = predict_essentia(float_musicnn_audio_buffer, embeddings_model, predictions_model)
    results_dict = dict(zip(MODEL_LABELS[model], np.round(np.mean(likelihoods, axis=0), 9)))

    # Aggregate results into merged categories
    probabilities = random_forrest_model.predict_proba(pd.DataFrame([results_dict]))
    genres_dict = {genre: prob for genre, prob in zip(random_forrest_model.classes_, probabilities.flatten())}

    # Display genres_dict using Streamlit
    maximum = max(results_dict, key=results_dict.get)
    text_output.markdown(f"**Top Tag:** {maximum, results_dict[maximum]}")
    json_output.write("Genres Dict:")
    json_output.json(genres_dict)

    # End timing
    end_process = time.time()
    print(f"> Inference done. Total processing time: {end_process - start_process:.4f}s")
    print(f"Genres Dict: {genres_dict}")
    print(f">===============================================================<")

async def process_audio(sound_chunk_queue, model, embeddings_model, predictions_model, random_forrest_model, text_output, json_output):
    sound_chunk = pydub.AudioSegment.empty()
    min_buffer_size = int(0.5 * SAMPLE_RATE)  # 0.5 seconds buffer

    while True:
        sound = await sound_chunk_queue.get()  # Get sound from the queue

        # Debugging print: track the queue size after consuming a chunk
        print(f"ProcessAudio: Retrieved frame from queue, current queue size: {sound_chunk_queue.qsize()}")
        sound_chunk += sound

        print(f"Current sound_chunk length: {len(sound_chunk)}, required: {min_buffer_size}")

        # If we have enough audio data (2 seconds), run inference
        if len(sound_chunk) >= min_buffer_size:
            print(f"Running inference with sound chunk of length {len(sound_chunk)}")
            await run_inference(sound_chunk, model, embeddings_model, predictions_model, random_forrest_model, text_output, json_output)

            # Keep only the last 0.5 seconds of audio for overlap in next inference
            last_frames = int(-0.5 * min_buffer_size)
            sound_chunk = sound_chunk[last_frames:]
        else:
            # Force inference every few cycles even if the buffer isn't quite full (if no new data is coming)
            if sound_chunk_queue.qsize() == 0 and len(sound_chunk) > int(1.5 * SAMPLE_RATE):  # Check buffer length, even with slightly smaller chunks
                print(f"Running forced inference with incomplete sound chunk of length {len(sound_chunk)}")
                await run_inference(sound_chunk, model, embeddings_model, predictions_model, random_forrest_model, text_output, json_output)
                sound_chunk = pydub.AudioSegment.empty()  # Clear buffer after inference


def app_sst(model):
    webrtc_ctx = webrtc_streamer(key="music-to-genre",
                                 mode=WebRtcMode.SENDONLY,
                                 audio_receiver_size=1024 * 8,
                                 rtc_configuration={"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]},
                                 media_stream_constraints={"video": False, "audio": True})
    
    status_indicator = st.empty()
    text_output = st.empty()
    json_output = st.empty()

    # Load model
    embeddings_model, predictions_model = load_essentia_model(model, MODEL_LABELS[model])
    random_forrest_model = joblib.load('pretrained-classifiers/random_forest_model.pkl')

    sound_chunk_queue = asyncio.Queue()

    async def run_tasks():
        audio_task = asyncio.create_task(audio_collector(webrtc_ctx, sound_chunk_queue))
        process_task = asyncio.create_task(process_audio(sound_chunk_queue, model, embeddings_model, predictions_model, random_forrest_model, text_output, json_output))
        
        # Run both tasks concurrently until `webrtc_ctx.state.playing` is False
        while webrtc_ctx.state.playing:
            await asyncio.sleep(0.1)  # Non-blocking sleep to keep checking the state

        # Gracefully cancel tasks when streaming stops
        audio_task.cancel()
        process_task.cancel()

    # Start event loop for tasks
    asyncio.run(run_tasks())


def main():
    st.set_page_config(layout="wide")
    st.header("Real Time Music Genre Classification")
    st.markdown("""This demo app is using musicnn as a recognition engine.""")
    sound_only_page = "Sound only demo"
    app_mode = st.selectbox("Choose the app mode", [sound_only_page])
    model = st.selectbox("Choose the model", ["ESSENTIA_genre_discogs400-discogs-effnet-1-random-forrest"])
    if app_mode == sound_only_page:
        app_sst(model)

if __name__ == "__main__":
    logging.basicConfig(format="[%(asctime)s] %(levelname)7s from %(name)s in %(pathname)s:%(lineno)d: %(message)s", force=True)
    logger.setLevel(level=logging.DEBUG if DEBUG else logging.INFO)

    st_webrtc_logger = logging.getLogger("streamlit_webrtc")
    st_webrtc_logger.setLevel(logging.DEBUG)

    fsevents_logger = logging.getLogger("fsevents")
    fsevents_logger.setLevel(logging.WARNING)

    main()

I would also welcome any ideas to improve this for speed and performance.