real_time_processor.py

#!/usr/bin/env python3
"""
Real-Time Processor - Unique Feature for Live Audio Transcription
Handles real-time audio processing, live transcription, and streaming capabilities
"""

import threading
import queue
import time
import numpy as np
from typing import Dict, Any, Optional, Callable
import logging

# Audio libraries
import pyaudio
import whisper
import webrtcvad
from scipy import signal

logger = logging.getLogger(__name__)

class RealTimeProcessor:
    """Real-time audio processing and transcription system"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.sample_rate = config['processing']['audio_sample_rate']
        self.whisper_model_name = config['models']['whisper_model']
        
        # Audio streaming parameters
        self.chunk_duration = 1.0  # seconds
        self.chunk_size = int(self.sample_rate * self.chunk_duration)
        self.overlap_duration = 0.3  # seconds overlap for continuity
        self.overlap_size = int(self.sample_rate * self.overlap_duration)
        
        # Initialize components
        self.whisper_model = whisper.load_model(self.whisper_model_name)
        self.vad = webrtcvad.Vad(2)  # Aggressiveness level 0-3
        
        # Streaming state
        self.is_streaming = False
        self.audio_queue = queue.Queue()
        self.transcription_queue = queue.Queue()
        self.audio_buffer = np.array([], dtype=np.float32)
        
        # Callbacks
        self.transcription_callback: Optional[Callable] = None
        self.voice_activity_callback: Optional[Callable] = None
        
    def start_streaming(self, 
                       transcription_callback: Callable[[str, Dict], None] = None,
                       voice_activity_callback: Callable[[bool], None] = None):
        """Start real-time audio streaming and transcription"""
        
        self.transcription_callback = transcription_callback
        self.voice_activity_callback = voice_activity_callback
        
        if self.is_streaming:
            logger.warning("Streaming already active")
            return
            
        self.is_streaming = True
        
        # Start audio capture thread
        self.audio_thread = threading.Thread(target=self._audio_capture_loop)
        self.audio_thread.daemon = True
        self.audio_thread.start()
        
        # Start transcription thread
        self.transcription_thread = threading.Thread(target=self._transcription_loop)
        self.transcription_thread.daemon = True
        self.transcription_thread.start()
        
        logger.info("Real-time processing started")
        
    def stop_streaming(self):
        """Stop real-time processing"""
        self.is_streaming = False
        
        if hasattr(self, 'audio_thread'):
            self.audio_thread.join(timeout=2.0)
        if hasattr(self, 'transcription_thread'):
            self.transcription_thread.join(timeout=2.0)
            
        logger.info("Real-time processing stopped")
        
    def _audio_capture_loop(self):
        """Capture audio from microphone in real-time"""
        try:
            # Initialize PyAudio
            audio = pyaudio.PyAudio()
            
            # Open audio stream
            stream = audio.open(
                format=pyaudio.paFloat32,
                channels=1,
                rate=self.sample_rate,
                input=True,
                frames_per_buffer=self.chunk_size,
                stream_callback=None
            )
            
            logger.info("Audio capture started")
            
            while self.is_streaming:
                # Read audio chunk
                try:
                    audio_data = stream.read(
                        self.chunk_size, 
                        exception_on_overflow=False
                    )
                    
                    # Convert to numpy array
                    audio_chunk = np.frombuffer(audio_data, dtype=np.float32)
                    
                    # Add to processing queue
                    self.audio_queue.put(audio_chunk)
                    
                except Exception as e:
                    logger.error(f"Audio capture error: {e}")
                    continue
                    
        except Exception as e:
            logger.error(f"Audio stream initialization failed: {e}")
        finally:
            try:
                stream.stop_stream()
                stream.close()
                audio.terminate()
            except:
                pass
                
    def _transcription_loop(self):
        """Process audio chunks for transcription"""
        
        while self.is_streaming:
            try:
                # Get audio chunk with timeout
                try:
                    audio_chunk = self.audio_queue.get(timeout=1.0)
                except queue.Empty:
                    continue
                    
                # Add to buffer with overlap handling
                self.audio_buffer = np.concatenate([self.audio_buffer, audio_chunk])
                
                # Process if buffer is large enough
                if len(self.audio_buffer) >= self.chunk_size + self.overlap_size:
                    # Extract processing chunk
                    process_chunk = self.audio_buffer[:self.chunk_size + self.overlap_size]
                    
                    # Keep overlap for next iteration
                    self.audio_buffer = self.audio_buffer[self.chunk_size:]
                    
                    # Voice Activity Detection
                    has_voice = self._detect_voice_activity(process_chunk)
                    
                    if self.voice_activity_callback:
                        self.voice_activity_callback(has_voice)
                        
                    # Transcribe if voice is detected
                    if has_voice:
                        transcription = self._transcribe_chunk(process_chunk)
                        
                        if transcription['text'].strip() and self.transcription_callback:
                            # Add timing information
                            transcription['timestamp'] = time.time()
                            transcription['chunk_duration'] = len(process_chunk) / self.sample_rate
                            
                            self.transcription_callback(transcription['text'], transcription)
                            
            except Exception as e:
                logger.error(f"Transcription loop error: {e}")
                continue
                
    def _detect_voice_activity(self, audio_chunk: np.ndarray) -> bool:
        """Detect voice activity in audio chunk"""
        try:
            # Convert to 16-bit PCM for WebRTC VAD
            audio_pcm = (audio_chunk * 32767).astype(np.int16)
            
            # WebRTC VAD requires specific frame sizes (10, 20, or 30 ms)
            frame_duration = 30  # ms
            frame_size = int(self.sample_rate * frame_duration / 1000)
            
            voice_frames = 0
            total_frames = 0
            
            # Process in frames
            for i in range(0, len(audio_pcm) - frame_size, frame_size):
                frame = audio_pcm[i:i + frame_size].tobytes()
                
                try:
                    is_speech = self.vad.is_speech(frame, self.sample_rate)
                    if is_speech:
                        voice_frames += 1
                    total_frames += 1
                except:
                    continue
                    
            # Return True if significant portion contains voice
            if total_frames > 0:
                voice_ratio = voice_frames / total_frames
                return voice_ratio > 0.3
                
        except Exception as e:
            logger.error(f"Voice activity detection failed: {e}")
            
        # Fallback to energy-based detection
        return self._energy_based_vad(audio_chunk)
        
    def _energy_based_vad(self, audio_chunk: np.ndarray) -> bool:
        """Simple energy-based voice activity detection"""
        try:
            # Calculate RMS energy
            rms_energy = np.sqrt(np.mean(audio_chunk**2))
            
            # Dynamic threshold based on recent history
            # This is a simplified implementation
            energy_threshold = 0.01  # Adjust based on testing
            
            return rms_energy > energy_threshold
            
        except Exception as e:
            logger.error(f"Energy-based VAD failed: {e}")
            return False
            
    def _transcribe_chunk(self, audio_chunk: np.ndarray) -> Dict[str, Any]:
        """Transcribe audio chunk using Whisper"""
        try:
            # Ensure audio is normalized
            if np.max(np.abs(audio_chunk)) > 1.0:
                audio_chunk = audio_chunk / np.max(np.abs(audio_chunk))
                
            # Transcribe with Whisper
            result = self.whisper_model.transcribe(
                audio_chunk,
                language=None,  # Auto-detect
                task="transcribe",
                fp16=False  # Use fp32 for better compatibility
            )
            
            return result
            
        except Exception as e:
            logger.error(f"Chunk transcription failed: {e}")
            return {"text": "", "language": "unknown"}
            
    def get_stream_statistics(self) -> Dict[str, Any]:
        """Get real-time processing statistics"""
        return {
            'is_streaming': self.is_streaming,
            'audio_queue_size': self.audio_queue.qsize(),
            'transcription_queue_size': self.transcription_queue.qsize(),
            'buffer_length_seconds': len(self.audio_buffer) / self.sample_rate,
            'sample_rate': self.sample_rate,
            'chunk_duration': self.chunk_duration,
            'overlap_duration': self.overlap_duration
        }
        
    def configure_sensitivity(self, 
                            vad_aggressiveness: int = 2,
                            energy_threshold: float = 0.01):
        """Configure voice activity detection sensitivity"""
        try:
            # Update VAD aggressiveness (0-3, higher = more aggressive)
            self.vad = webrtcvad.Vad(vad_aggressiveness)
            
            # Store energy threshold for fallback VAD
            self.energy_threshold = energy_threshold
            
            logger.info(f"VAD sensitivity updated: aggressiveness={vad_aggressiveness}")
            
        except Exception as e:
            logger.error(f"Failed to configure sensitivity: {e}")
            
    def process_audio_file_realtime(self, 
                                   file_path: str, 
                                   playback_speed: float = 1.0):
        """Simulate real-time processing of an audio file"""
        try:
            import librosa
            
            # Load audio file
            audio_data, sr = librosa.load(file_path, sr=self.sample_rate)
            
            # Calculate chunk timing
            chunk_samples = int(self.sample_rate * self.chunk_duration)
            delay_between_chunks = self.chunk_duration / playback_speed
            
            logger.info(f"Starting real-time simulation of {file_path}")
            
            # Process in chunks with realistic timing
            for i in range(0, len(audio_data), chunk_samples):
                if not self.is_streaming:
                    break
                    
                chunk = audio_data[i:i + chunk_samples]
                
                # Pad if chunk is too short
                if len(chunk) < chunk_samples:
                    chunk = np.pad(chunk, (0, chunk_samples - len(chunk)))
                    
                # Voice activity detection
                has_voice = self._detect_voice_activity(chunk)
                
                if self.voice_activity_callback:
                    self.voice_activity_callback(has_voice)
                    
                # Transcribe if voice detected
                if has_voice:
                    transcription = self._transcribe_chunk(chunk)
                    
                    if transcription['text'].strip() and self.transcription_callback:
                        transcription['timestamp'] = time.time()
                        transcription['file_position'] = i / sr
                        transcription['chunk_duration'] = len(chunk) / self.sample_rate
                        
                        self.transcription_callback(transcription['text'], transcription)
                        
                # Wait for next chunk (simulate real-time)
                time.sleep(delay_between_chunks)
                
        except Exception as e:
            logger.error(f"Real-time file processing failed: {e}")


class LiveTranscriptionManager:
    """Manages live transcription sessions with text smoothing and correction"""
    
    def __init__(self, config: Dict[str, Any]):
        self.config = config
        self.current_session = None
        self.session_history = []
        
    def start_session(self, session_name: str = None) -> str:
        """Start a new live transcription session"""
        session_id = session_name or f"session_{int(time.time())}"
        
        self.current_session = {
            'id': session_id,
            'start_time': time.time(),
            'transcripts': [],
            'raw_chunks': [],
            'processed_text': "",
            'confidence_scores': [],
            'language_detected': None
        }
        
        logger.info(f"Started transcription session: {session_id}")
        return session_id
        
    def add_transcription_chunk(self, text: str, metadata: Dict[str, Any]):
        """Add transcription chunk to current session"""
        if not self.current_session:
            self.start_session()
            
        chunk = {
            'text': text,
            'timestamp': time.time(),
            'metadata': metadata
        }
        
        self.current_session['transcripts'].append(chunk)
        self.current_session['raw_chunks'].append(text)
        
        # Update processed text with smoothing
        self._update_processed_text(text)
        
    def _update_processed_text(self, new_text: str):
        """Update processed text with smoothing and correction"""
        # Simple text smoothing - can be enhanced with more sophisticated NLP
        
        # Remove duplicate consecutive words/phrases
        words = new_text.split()
        if len(words) > 0:
            # Check for overlap with previous text
            current_words = self.current_session['processed_text'].split()
            
            # Find overlap and append only new words
            overlap_found = False
            for i in range(min(len(current_words), len(words))):
                if current_words[-i-1:] == words[:i+1]:
                    # Found overlap, append remaining words
                    new_words = words[i+1:]
                    if new_words:
                        self.current_session['processed_text'] += " " + " ".join(new_words)
                    overlap_found = True
                    break
                    
            if not overlap_found:
                # No overlap found, append with separator
                if self.current_session['processed_text']:
                    self.current_session['processed_text'] += " " + new_text
                else:
                    self.current_session['processed_text'] = new_text
                    
    def get_current_transcript(self) -> str:
        """Get current processed transcript"""
        if self.current_session:
            return self.current_session['processed_text']
        return ""
        
    def end_session(self) -> Dict[str, Any]:
        """End current transcription session and return summary"""
        if not self.current_session:
            return {}
            
        self.current_session['end_time'] = time.time()
        self.current_session['duration'] = self.current_session['end_time'] - self.current_session['start_time']
        
        # Calculate session statistics
        total_chunks = len(self.current_session['transcripts'])
        total_words = len(self.current_session['processed_text'].split())
        
        session_summary = {
            'session_id': self.current_session['id'],
            'duration': self.current_session['duration'],
            'total_chunks': total_chunks,
            'total_words': total_words,
            'words_per_minute': total_words / (self.current_session['duration'] / 60) if self.current_session['duration'] > 0 else 0,
            'final_transcript': self.current_session['processed_text'],
            'language_detected': self.current_session.get('language_detected', 'unknown')
        }
        
        # Archive session
        self.session_history.append(self.current_session.copy())
        self.current_session = None
        
        logger.info(f"Ended transcription session: {session_summary['session_id']}")
        return session_summary