Spaces:
Sleeping
Sleeping
| """ | |
| This script implements a Voice Activity Detection (VAD) system using a Convolutional Neural Network (CNN) model. | |
| The script uses Gradio for the user interface, NumPy and SciPy for signal | |
| processing, and ONNX Runtime for model inference. | |
| Modules: | |
| - gradio: For creating the web interface. | |
| - numpy: For numerical operations. | |
| - scipy: For signal processing. | |
| - onnxruntime: For running the ONNX model. | |
| - threading: For running model inference in a separate thread. | |
| - queue: For thread-safe communication between threads. | |
| """ | |
| import atexit | |
| import queue | |
| import threading | |
| import time | |
| import gradio as gr | |
| import numpy as np | |
| import onnxruntime as ort | |
| from scipy import signal | |
| from scipy.fft import rfft | |
| # Parameters | |
| ORIGINAL_FS = 48_000 # Original sampling frequency in Hz | |
| DECIMATION_FACTOR = 3 # Decimation factor | |
| TARGET_FS = ORIGINAL_FS // DECIMATION_FACTOR # Target sampling frequency in Hz | |
| FRAME_DURATION_MS = 25 # Frame duration in milliseconds | |
| FRAME_SIZE = int(ORIGINAL_FS * FRAME_DURATION_MS / 1000) # Frame size | |
| FRAME_STEP = FRAME_SIZE // 2 # 50% overlap | |
| # FFT parameters | |
| DECIMATED_FRAME_SIZE = FRAME_SIZE // DECIMATION_FACTOR | |
| FFT_SIZE = 1 | |
| while FFT_SIZE < DECIMATED_FRAME_SIZE: | |
| FFT_SIZE <<= 1 | |
| # Mel spectrogram parameters | |
| N_FILTERS = 40 # Number of mel filters | |
| FREQ_LOW = 300 # Lower frequency bound | |
| FREQ_HIGH = 8000 # Upper frequency bound | |
| # Model evaluation parameters | |
| INITIAL_FRAMES = 40 # Initial frames required for first model evaluation | |
| FRAMES_PER_EVAL = 20 # Run model every this many new frames after initial | |
| # EWMA parameter | |
| ALPHA = 0.9 # Weight for current observation in EWMA | |
| # Filter coefficients | |
| FILTER_COEFFICIENTS = np.array([0.0625, 0.125, 0.25, 0.5, 0.25, 0.125, 0.0625]) | |
| # Create window function (Hanning window) | |
| window = np.hanning(DECIMATED_FRAME_SIZE) | |
| # Create mel filterbank matrix | |
| def create_mel_filterbank(n_filt, freq_low, freq_high, n_fft, fs): | |
| """ | |
| Create a mel filterbank matrix for computing the mel spectrogram. | |
| Args: | |
| n_filt: Number of mel filters | |
| freq_low: Lower frequency bound | |
| freq_high: Upper frequency bound | |
| n_fft: FFT size | |
| fs: Sampling frequency | |
| Returns: | |
| filterbank: Mel filterbank matrix of shape (n_fft//2+1, n_filt) | |
| """ | |
| # Convert Hz to mel | |
| def hz_to_mel(hz): | |
| return 1125 * np.log(1 + hz / 700) | |
| # Convert mel to Hz | |
| def mel_to_hz(mel): | |
| return 700 * (np.exp(mel / 1125) - 1) | |
| # Compute points evenly spaced in mel scale | |
| lower_mel = hz_to_mel(freq_low) | |
| higher_mel = hz_to_mel(freq_high) | |
| mel_points = np.linspace(lower_mel, higher_mel, n_filt + 2) | |
| # Convert mel points to Hz | |
| hz_points = mel_to_hz(mel_points) | |
| # Convert Hz points to FFT bin indices | |
| bin_indices = np.floor((n_fft + 1) * hz_points / fs).astype(int) | |
| # Create filterbank matrix | |
| filterbank = np.zeros((n_fft // 2 + 1, n_filt)) | |
| for i in range(n_filt): | |
| # For each filter, create a triangular filter | |
| for j in range(bin_indices[i], bin_indices[i + 1]): | |
| filterbank[j, i] = (j - bin_indices[i]) / ( | |
| bin_indices[i + 1] - bin_indices[i] | |
| ) | |
| for j in range(bin_indices[i + 1], bin_indices[i + 2]): | |
| filterbank[j, i] = (bin_indices[i + 2] - j) / ( | |
| bin_indices[i + 2] - bin_indices[i + 1] | |
| ) | |
| return filterbank | |
| # Create mel filterbank | |
| mel_filterbank = create_mel_filterbank( | |
| N_FILTERS, FREQ_LOW, FREQ_HIGH, FFT_SIZE, TARGET_FS | |
| ) | |
| # Initialize mel spectrogram image | |
| mel_spectrogram_image = np.zeros((N_FILTERS, N_FILTERS), dtype=np.float32) | |
| # Load the ONNX model | |
| session_options = ort.SessionOptions() | |
| session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL | |
| session_options.intra_op_num_threads = 2 | |
| session = ort.InferenceSession("./model/cnn-vad.onnx", session_options) | |
| input_name = session.get_inputs()[0].name | |
| output_name = session.get_outputs()[0].name | |
| # Create thread-safe queues for ML tasks and results | |
| inference_queue = queue.Queue(maxsize=10) # Queue for pending inference tasks | |
| result_queue = queue.Queue(maxsize=10) # Queue for inference results | |
| stop_event = threading.Event() # Event to signal thread termination | |
| # ML worker thread function | |
| def ml_worker(): | |
| """Worker thread function for running ML inference.""" | |
| while not stop_event.is_set(): | |
| try: | |
| # Try to get a task with timeout to allow checking stop_event periodically | |
| mel_image = inference_queue.get(timeout=0.1) | |
| # Run inference | |
| outputs = session.run( | |
| [output_name], | |
| { | |
| input_name: mel_image.astype(np.float32).reshape( | |
| -1, N_FILTERS, N_FILTERS | |
| ) | |
| }, | |
| )[0] | |
| # Get speech probability score and put in result queue | |
| current_score = outputs[0][1] | |
| result_queue.put(current_score) | |
| # Mark task as done | |
| inference_queue.task_done() | |
| except queue.Empty: | |
| # No pending tasks, just continue | |
| continue | |
| except Exception as e: | |
| print(f"Error in ML worker thread: {e}") | |
| # Still mark task as done if there was one | |
| if not inference_queue.empty(): | |
| inference_queue.task_done() | |
| # Start ML worker thread | |
| ml_thread = threading.Thread(target=ml_worker, daemon=True) | |
| ml_thread.start() | |
| # Initialize stream state | |
| initial_state = { | |
| "buffer": np.zeros(0, dtype=np.float32), # Audio buffer | |
| "frames_processed": 0, # Number of frames processed | |
| "ewma_score": 0.0, # EWMA of speech detection score | |
| "mel_image": np.zeros( | |
| (N_FILTERS, N_FILTERS), dtype=np.float32 | |
| ), # Mel spectrogram image | |
| "inference_pending": False, # Flag to track if inference is pending | |
| "last_inference_time": 0, # Timestamp of last inference | |
| } | |
| def resample_audio(audio, orig_sr, target_sr): | |
| """ | |
| Resample audio from original sample rate to target sample rate. | |
| Simple implementation using linear interpolation. | |
| Args: | |
| audio: Audio data | |
| orig_sr: Original sample rate | |
| target_sr: Target sample rate | |
| Returns: | |
| Resampled audio | |
| """ | |
| if orig_sr == target_sr: | |
| return audio | |
| # For more accurate resampling in production, consider using: | |
| # from scipy import signal | |
| return signal.resample_poly(audio, target_sr, orig_sr) | |
| # # Simple linear interpolation (less accurate but faster) | |
| # duration = len(audio) / orig_sr | |
| # new_length = int(duration * target_sr) | |
| # indices = np.linspace(0, len(audio) - 1, new_length) | |
| # indices = indices.astype(np.int32) | |
| # return audio[indices] | |
| def process_frame(frame, mel_image): | |
| """ | |
| Process a single frame of audio data. | |
| Args: | |
| frame: Audio frame data | |
| mel_image: Current mel spectrogram image | |
| Returns: | |
| Updated mel spectrogram image | |
| """ | |
| # 1. Apply FIR filter | |
| filtered_frame = signal.lfilter(FILTER_COEFFICIENTS, [1.0], frame) | |
| # Clamp values | |
| filtered_frame = np.clip(filtered_frame, -1.0, 1.0) | |
| # 2. Decimate the frame | |
| decimated_frame = filtered_frame[::DECIMATION_FACTOR] | |
| # Apply window and pad with zeros to FFT_SIZE | |
| padded_frame = np.zeros(FFT_SIZE) | |
| padded_frame[:DECIMATED_FRAME_SIZE] = decimated_frame * window | |
| # 3. Perform FFT | |
| fft_result = rfft(padded_frame) | |
| # Compute power spectrum | |
| power_spectrum = np.abs(fft_result) ** 2 | |
| # 4. Calculate mel spectrogram | |
| mel_power = np.dot(power_spectrum[: FFT_SIZE // 2 + 1], mel_filterbank) | |
| # Log mel spectrogram | |
| mel_power = np.log(mel_power + 1e-8) | |
| # Update mel spectrogram image (shift up and add new row at bottom) | |
| mel_image = np.roll(mel_image, -1, axis=0) | |
| mel_image[-1] = mel_power | |
| return mel_image | |
| def detect(state, new_chunk): | |
| """ | |
| Detects speech in an audio stream using a Voice Activity Detection (VAD) model. | |
| Args: | |
| state: Current state dictionary containing: | |
| - buffer: Audio buffer | |
| - frames_processed: Number of frames processed so far | |
| - ewma_score: Exponentially weighted moving average of speech detection score | |
| - mel_image: Current mel spectrogram image | |
| - inference_pending: Flag indicating if an inference is pending | |
| - last_inference_time: Timestamp of last inference | |
| new_chunk: A tuple containing the sample rate (sr) and the audio data (y). | |
| Returns: | |
| A tuple containing the updated state, a string indicating whether speech was detected, | |
| and an HTML element for the visual indicator. | |
| """ | |
| # Initialize state if it's the first call | |
| if state is None: | |
| state = initial_state.copy() | |
| state["mel_image"] = np.copy(initial_state["mel_image"]) | |
| sr, y = new_chunk | |
| # If no audio data, return current state | |
| if y is None or len(y) == 0: | |
| return ( | |
| state, | |
| f"No audio detected (Score: {state['ewma_score']:.2f})", | |
| generate_indicator(state["ewma_score"]), | |
| ) | |
| # Pre-processing: Convert to mono if stereo | |
| if len(y.shape) > 1: | |
| y = np.mean(y, axis=1) | |
| # Convert to float32 [-1.0, 1.0] | |
| y = y.astype(np.float32) / 32768 | |
| # Resample if necessary | |
| if sr != ORIGINAL_FS: | |
| y = resample_audio(y, sr, ORIGINAL_FS) | |
| # Append new audio to buffer | |
| buffer = np.concatenate([state["buffer"], y]) | |
| # Process frames from buffer | |
| new_frames_processed = 0 | |
| run_inference = False | |
| current_time = time.time() | |
| # Process as many complete frames as possible | |
| while len(buffer) >= FRAME_SIZE: | |
| # Extract frame | |
| frame = buffer[:FRAME_SIZE] | |
| buffer = buffer[FRAME_STEP:] # Advance by hop size (50% overlap) | |
| # Process frame and update mel spectrogram image | |
| state["mel_image"] = process_frame(frame, state["mel_image"]) | |
| new_frames_processed += 1 | |
| state["frames_processed"] += 1 | |
| # Determine if we should run inference | |
| if ( | |
| state["frames_processed"] >= INITIAL_FRAMES | |
| and (state["frames_processed"] - INITIAL_FRAMES) % FRAMES_PER_EVAL == 0 | |
| ): | |
| run_inference = True | |
| # Update buffer in state | |
| state["buffer"] = buffer | |
| # Check for completed inference results | |
| if not result_queue.empty(): | |
| # Get the result from the queue | |
| current_score = result_queue.get() | |
| # Update EWMA | |
| if state["ewma_score"] == 0: # First evaluation | |
| state["ewma_score"] = current_score | |
| else: | |
| state["ewma_score"] = ( | |
| ALPHA * current_score + (1 - ALPHA) * state["ewma_score"] | |
| ) | |
| # Mark that we're no longer waiting for inference | |
| state["inference_pending"] = False | |
| # Run VAD model if criteria are met and no inference is currently pending | |
| if ( | |
| ( | |
| run_inference | |
| or ( | |
| state["frames_processed"] >= INITIAL_FRAMES | |
| and new_frames_processed > 0 | |
| and current_time - state["last_inference_time"] > 0.1 # Rate limiting | |
| ) | |
| ) | |
| and not state["inference_pending"] | |
| and not inference_queue.full() | |
| ): | |
| # Instead of running inference here, queue it for the worker thread | |
| inference_queue.put(np.copy(state["mel_image"])) | |
| state["inference_pending"] = True | |
| state["last_inference_time"] = current_time | |
| # Determine result based on EWMA score | |
| if state["frames_processed"] < INITIAL_FRAMES: | |
| message = ( | |
| f"Building audio context... ({state['frames_processed']}/{INITIAL_FRAMES})" | |
| ) | |
| indicator = generate_indicator(0) # Inactive during initialization | |
| elif state["ewma_score"] > 0.5: | |
| message = f"Speech Detected (Score: {state['ewma_score']:.2f})" | |
| indicator = generate_indicator(state["ewma_score"]) | |
| else: | |
| message = f"No Speech Detected (Score: {state['ewma_score']:.2f})" | |
| indicator = generate_indicator(state["ewma_score"]) | |
| return state, message, indicator | |
| def generate_indicator(score): | |
| """ | |
| Generate an HTML indicator that lights up based on the speech detection score. | |
| Args: | |
| score: Speech detection score (0.0 to 1.0) | |
| Returns: | |
| HTML string for the visual indicator | |
| """ | |
| # Set colors based on score | |
| if score > 0.5: | |
| # Active - green glow | |
| brightness = min(100, 50 + int(score * 50)) # 50-100% brightness based on score | |
| color = f"rgba(144, 219, 130, {score:.1f})" # Green with opacity based on score | |
| glow = f"0 0 {int(score * 20)}px rgba(0, 255, 0, {score:.1f})" # Glow effect | |
| else: | |
| # Inactive - dim red | |
| brightness = max(10, int(score * 50)) # 0-25% brightness based on score | |
| color = f"rgba(255, 0, 0, {max(0.1, score):.1f})" # Red with low opacity | |
| glow = "none" # No glow when inactive | |
| # Create HTML for the indicator | |
| html = f""" | |
| <div style="display: flex; justify-content: center; margin-top: 10px; margin-bottom: 10px;"> | |
| <div style=" | |
| width: 80px; | |
| height: 80px; | |
| border-radius: 50%; | |
| background-color: {color}; | |
| box-shadow: {glow}; | |
| transition: all 0.3s ease; | |
| display: flex; | |
| justify-content: center; | |
| align-items: center; | |
| font-weight: bold; | |
| color: white; | |
| filter: brightness({brightness}%); | |
| "> | |
| {int(score * 100)}% | |
| </div> | |
| </div> | |
| """ | |
| return html | |
| # Create the Gradio interface | |
| with gr.Blocks(theme=gr.themes.Ocean()) as demo: | |
| gr.Markdown("# Voice Activity Detection") | |
| gr.Markdown( | |
| "Speak into your microphone to see the voice activity detection in action." | |
| ) | |
| # State for maintaining app state between calls | |
| state = gr.State(None) | |
| # Audio input | |
| audio_input = gr.Audio( | |
| sources=["microphone"], | |
| streaming=True, | |
| autoplay=True, | |
| elem_id="mic_input", | |
| ) | |
| # Visual indicator and text output | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| indicator_html = gr.HTML(generate_indicator(0)) | |
| with gr.Column(scale=2): | |
| text_output = gr.Textbox(label="Detection Result") | |
| # Set up the processing function | |
| audio_input.stream( | |
| detect, | |
| inputs=[state, audio_input], | |
| outputs=[state, text_output, indicator_html], | |
| show_progress=False, | |
| ) | |
| gr.Markdown(""" | |
| ## How it works | |
| This app uses a Convolutional Neural Network (CNN) to detect speech in audio. | |
| - The indicator lights up **green** when speech is detected | |
| - The indicator turns **red** when no speech is detected | |
| - The percentage shows the confidence level of speech detection | |
| - The model and processing is listed in [this paper](https://ieeexplore.ieee.org/abstract/document/8278160) | |
| """) | |
| # Cleanup function to be called when the app is closed | |
| def cleanup(): | |
| # Signal the worker thread to stop | |
| stop_event.set() | |
| # Wait for the thread to finish (with timeout) | |
| ml_thread.join(timeout=1.0) | |
| print("ML worker thread stopped") | |
| # Register cleanup handler | |
| atexit.register(cleanup) | |
| # Launch the interface | |
| demo.launch() | |