SussurroXRest / utils /audio_buffer.py
LucaR84's picture
code refactor
182c2c2
import numpy as np
import soundfile as sf
import librosa
from io import BytesIO
from typing import Optional
class AudioBuffer:
"""Buffer to accumulate streaming audio data"""
def __init__(self, chunk_duration: float = 30.0, sample_rate: int = 16000):
self.chunk_duration = chunk_duration
self.sample_rate = sample_rate
self.buffer = np.array([], dtype=np.float32)
self.target_chunk_size = int(sample_rate * chunk_duration)
def add_data(self, audio_data: bytes) -> Optional[np.ndarray]:
"""Add audio data to buffer and return chunk if ready"""
# Convert bytes to audio array
audio_buffer = BytesIO(audio_data)
audio_chunk, sample_rate = sf.read(audio_buffer)
# Ensure mono audio
if len(audio_chunk.shape) > 1:
audio_chunk = audio_chunk.mean(axis=1)
# Resample to target sample rate if needed
if sample_rate != self.sample_rate:
audio_chunk = librosa.resample(audio_chunk, orig_sr=sample_rate, target_sr=self.sample_rate)
self.buffer = np.concatenate([self.buffer, audio_chunk.astype(np.float32)])
if len(self.buffer) >= self.target_chunk_size:
# Extract chunk
chunk = self.buffer[:self.target_chunk_size]
# Keep remaining data in buffer
self.buffer = self.buffer[self.target_chunk_size:]
return chunk
return None
def get_remaining(self) -> np.ndarray:
"""Get any remaining audio data"""
return self.buffer