dudong-v2/workers/audio_worker.py

"""
Audio Worker Module

Worker thread for async audio processing and ripeness classification.
"""

from typing import Optional
import logging

from PyQt5.QtCore import pyqtSignal
from PyQt5.QtGui import QPixmap

from workers.base_worker import BaseWorker, WorkerSignals
from models.audio_model import AudioModel

logger = logging.getLogger(__name__)


class AudioWorkerSignals(WorkerSignals):
    """
    Signals specific to audio processing.
    
    Signals:
        result_ready: Emitted when prediction is complete
            (waveform: QPixmap, spectrogram: QPixmap, class_name: str, 
             confidence: float, probabilities: dict, knock_count: int)
    """
    result_ready = pyqtSignal(QPixmap, QPixmap, str, float, dict, int)


class AudioWorker(BaseWorker):
    """
    Worker for processing audio files and predicting ripeness.

    Runs AudioModel inference in a background thread without blocking UI.
    Supports multiple audio formats (WAV, MP3, FLAC, OGG, M4A, AAC, WMA).

    Attributes:
        audio_path: Path to audio file (multiple formats supported)
        model: AudioModel instance
        signals: AudioWorkerSignals for emitting results
    """
    
    def __init__(self, audio_path: str, model: Optional[AudioModel] = None):
        """
        Initialize the audio worker.

        Args:
            audio_path: Path to audio file to process (supports multiple formats)
            model: AudioModel instance (if None, creates new one)
        """
        super().__init__()
        self.audio_path = audio_path
        self.model = model
        
        # Replace base signals with audio-specific signals
        self.signals = AudioWorkerSignals()
        
        # If no model provided, create and load one
        if self.model is None:
            self.model = AudioModel(device='cpu')  # TensorFlow works best on CPU
            self.model.load()
        
        logger.info(f"AudioWorker created for: {audio_path}")
    
    def process(self):
        """
        Process the audio file and predict ripeness.
        
        Emits result_ready signal with prediction results.
        """
        if self.is_cancelled():
            logger.info("AudioWorker cancelled before processing")
            return
        
        # Update progress
        self.emit_progress(10, "Loading audio file...")
        
        if not self.model.is_loaded:
            logger.warning("⚠️  Model not loaded, loading now...")
            self.emit_progress(30, "Loading model...")
            if not self.model.load():
                logger.error("❌ Failed to load audio model")
                raise RuntimeError("Failed to load audio model")
            logger.info("✓ Model loaded successfully")
        
        # Process audio
        self.emit_progress(50, "Detecting knocks and generating visualizations...")
        
        logger.info(f"Processing audio: {self.audio_path}")
        result = self.model.predict(self.audio_path)
        logger.info(f"Prediction result success: {result.get('success')}")
        
        if self.is_cancelled():
            logger.info("AudioWorker cancelled during processing")
            return
        
        if not result['success']:
            logger.error(f"❌ Prediction failed: {result.get('error')}")
            raise RuntimeError(result['error'])
        
        logger.info(f"✓ Prediction successful: {result.get('class_name')} ({result.get('confidence'):.2%})")
        
        self.emit_progress(90, "Finalizing results...")
        
        # Emit results (handle None images gracefully with placeholders)
        waveform_image = result.get('waveform_image')
        spectrogram_image = result.get('spectrogram_image')
        
        if waveform_image is None:
            # Create a placeholder pixmap for missing waveform
            from PyQt5.QtGui import QPixmap, QColor
            waveform_image = QPixmap(400, 200)
            waveform_image.fill(QColor("#2c3e50"))
        
        if spectrogram_image is None:
            # Create a placeholder pixmap for missing spectrogram
            from PyQt5.QtGui import QPixmap, QColor
            spectrogram_image = QPixmap(400, 200)
            spectrogram_image.fill(QColor("#2c3e50"))

        self.signals.result_ready.emit(
            waveform_image,
            spectrogram_image,
            result['class_name'],
            result['confidence'],
            result['probabilities'],
            result.get('knock_count', 0)
        )
        
        self.emit_progress(100, "Complete!")
        
        logger.info(f"AudioWorker completed: {result['class_name']} ({result['confidence']:.2%}) - {result.get('knock_count', 0)} knocks")