data_processing/audio_processor.py

"""Audio processing module for Myanmar Ghost project."""

import logging
from pathlib import Path
from typing import Optional, Tuple

import librosa
import numpy as np
import soundfile as sf
from scipy.signal import butter, filtfilt

logger = logging.getLogger(__name__)


class AudioProcessor:
    """Process audio files for Myanmar speech recognition."""

    def __init__(
        self,
        sample_rate: int = 16000,
        n_fft: int = 512,
        hop_length: int = 160,
        n_mels: int = 80,
    ):
        self.sample_rate = sample_rate
        self.n_fft = n_fft
        self.hop_length = hop_length
        self.n_mels = n_mels

    def load_audio(self, path: str) -> Tuple[np.ndarray, int]:
        """Load audio file and resample to target sample rate."""
        audio, sr = librosa.load(path, sr=self.sample_rate)
        logger.info(f"Loaded audio from {path}: {len(audio)} samples at {sr}Hz")
        return audio, sr

    def normalize_audio(self, audio: np.ndarray) -> np.ndarray:
        """Normalize audio to [-1, 1] range."""
        max_val = np.abs(audio).max()
        if max_val > 0:
            audio = audio / max_val
        return audio

    def remove_silence(
        self,
        audio: np.ndarray,
        threshold_db: float = -40,
        min_silence_duration: float = 0.3,
    ) -> np.ndarray:
        """Remove silence from audio based on energy threshold."""
        intervals = librosa.effects.split(
            audio,
            top_db=-threshold_db,
            frame_length=self.n_fft,
            hop_length=self.hop_length,
        )
        
        if len(intervals) == 0:
            return audio
        
        min_samples = int(min_silence_duration * self.sample_rate)
        non_silent = []
        
        for start, end in intervals:
            if end - start >= min_samples:
                non_silent.append(audio[start:end])
        
        if non_silent:
            return np.concatenate(non_silent)
        return audio

    def apply_bandpass_filter(
        self,
        audio: np.ndarray,
        low_freq: float = 80,
        high_freq: float = 7500,
    ) -> np.ndarray:
        """Apply bandpass filter to focus on speech frequencies."""
        nyquist = self.sample_rate / 2
        low = low_freq / nyquist
        high = high_freq / nyquist
        
        if low < 0:
            low = 0.001
        if high > 1:
            high = 0.999
            
        b, a = butter(4, [low, high], btype="band")
        filtered = filtfilt(b, a, audio)
        return filtered

    def reduce_noise(
        self,
        audio: np.ndarray,
        noise_profile: Optional[np.ndarray] = None,
    ) -> np.ndarray:
        """Reduce background noise using spectral subtraction."""
        if noise_profile is None:
            noise_profile = audio[: int(0.1 * self.sample_rate)]
        
        noise_spectrum = np.abs(np.fft.rfft(noise_profile))
        noise_magnitude = np.mean(noise_spectrum, axis=0)
        
        audio_spectrum = np.abs(np.fft.rfft(audio))
        cleaned = np.maximum(
            audio_spectrum - noise_magnitude[:, None],
            audio_spectrum * 0.1,
        )
        cleaned = cleaned * np.exp(1j * np.fft.rfft(audio).angle())
        
        return np.fft.irfft(cleaned)

    def extract_mel_spectrogram(self, audio: np.ndarray) -> np.ndarray:
        """Extract mel spectrogram features."""
        mel_spec = librosa.feature.melspectrogram(
            y=audio,
            sr=self.sample_rate,
            n_fft=self.n_fft,
            hop_length=self.hop_length,
            n_mels=self.n_mels,
        )
        log_mel = librosa.power_to_db(mel_spec, ref=np.max)
        return log_mel

    def extract_prosody_features(self, audio: np.ndarray) -> dict:
        """Extract prosodic features (pitch, energy, speaking rate)."""
        pitches, magnitudes = librosa.piptrack(
            y=audio,
            sr=self.sample_rate,
            n_fft=self.n_fft,
            hop_length=self.hop_length,
        )
        
        pitch_values = []
        for i in range(pitches.shape[1]):
            index = magnitudes[:, i].argmax()
            pitch = pitches[index, i]
            if pitch > 0:
                pitch_values.append(pitch)
        
        rms = librosa.feature.rms(y=audio, hop_length=self.hop_length)[0]
        
        return {
            "mean_pitch": np.mean(pitch_values) if pitch_values else 0,
            "pitch_std": np.std(pitch_values) if pitch_values else 0,
            "pitch_range": (np.min(pitch_values) if pitch_values else 0,
                          np.max(pitch_values) if pitch_values else 0),
            "mean_energy": np.mean(rms),
            "energy_std": np.std(rms),
        }

    def process_file(
        self,
        input_path: str,
        output_path: str,
        remove_silence: bool = True,
        apply_filter: bool = True,
    ) -> dict:
        """Process a single audio file."""
        audio, sr = self.load_audio(input_path)
        audio = self.normalize_audio(audio)
        
        if apply_filter:
            audio = self.apply_bandpass_filter(audio)
        
        if remove_silence:
            audio = self.remove_silence(audio)
        
        prosody = self.extract_prosody_features(audio)
        
        sf.write(output_path, audio, self.sample_rate)
        logger.info(f"Saved processed audio to {output_path}")
        
        return {
            "input_path": input_path,
            "output_path": output_path,
            "duration": len(audio) / self.sample_rate,
            "prosody": prosody,
        }

    def batch_process(
        self,
        input_dir: str,
        output_dir: str,
        pattern: str = "*.wav",
    ) -> list:
        """Process all audio files in a directory."""
        input_path = Path(input_dir)
        output_path = Path(output_dir)
        output_path.mkdir(parents=True, exist_ok=True)
        
        results = []
        for file_path in input_path.glob(pattern):
            out_file = output_path / file_path.name
            result = self.process_file(str(file_path), str(out_file))
            results.append(result)
        
        return results


def create_processor(config: dict = None) -> AudioProcessor:
    """Factory function to create AudioProcessor from config."""
    if config is None:
        config = {}
    
    return AudioProcessor(
        sample_rate=config.get("sample_rate", 16000),
        n_fft=config.get("n_fft", 512),
        hop_length=config.get("hop_length", 160),
        n_mels=config.get("n_mels", 80),
    )


if __name__ == "__main__":
    processor = create_processor()
    print("AudioProcessor initialized successfully")