Skip to content

how to test custom wake word after training the model? #13

Description

@tejas-p-28
import numpy as np
import onnxruntime as ort
import sounddevice as sd
import queue
import sys
from collections import deque
import librosa

# Configuration
MODEL_PATH = "/home/tejas/wakeword/arcosoph_A_v1.onnx"  # Update with your model path
SAMPLE_RATE = 16000
CHUNK_DURATION = 0.1  # 100ms chunks
CHUNK_SIZE = int(SAMPLE_RATE * CHUNK_DURATION)
THRESHOLD = 0.5  # Detection threshold (adjust based on your model)

# Mel spectrogram parameters for nanowakeword
N_MELS = 16  # Number of mel bins (from model input shape)
N_FRAMES = 96  # Number of time frames (from model input shape)
HOP_LENGTH = 160  # ~10ms hop
N_FFT = 400  # ~25ms window

class WakeWordDetector:
    def __init__(self, model_path, threshold=0.5):
        self.threshold = threshold
        self.session = ort.InferenceSession(model_path)
        
        # Get model input details
        self.input_name = self.session.get_inputs()[0].name
        self.input_shape = self.session.get_inputs()[0].shape
        
        print(f"Model loaded: {model_path}")
        print(f"Input name: {self.input_name}")
        print(f"Input shape: {self.input_shape}")
        
        # Calculate required audio samples
        self.required_samples = (N_FRAMES - 1) * HOP_LENGTH + N_FFT
        print(f"Required audio samples: {self.required_samples} ({self.required_samples/SAMPLE_RATE:.2f}s)")
        
        # Audio buffer - store raw audio samples
        self.audio_buffer = deque(maxlen=self.required_samples * 2)
        self.audio_queue = queue.Queue()
        
    def audio_callback(self, indata, frames, time, status):
        """Callback for audio stream"""
        if status:
            print(f"Audio status: {status}", file=sys.stderr)
        self.audio_queue.put(indata.copy())
    
    def audio_to_mel_spectrogram(self, audio):
        """Convert audio to mel spectrogram"""
        # Ensure audio is 1D float32
        audio = audio.flatten().astype(np.float32)
        
        # Normalize to [-1, 1] if needed
        if audio.max() > 1.0 or audio.min() < -1.0:
            audio = audio / 32768.0
        
        # Debug: Check audio stats
        if hasattr(self, 'debug_count') and self.debug_count < 3:
            print(f"\nDEBUG - Audio stats:")
            print(f"  Shape: {audio.shape}")
            print(f"  Min: {audio.min():.4f}, Max: {audio.max():.4f}")
            print(f"  Mean: {audio.mean():.4f}, Std: {audio.std():.4f}")
            print(f"  RMS: {np.sqrt(np.mean(audio**2)):.4f}")
        
        # Compute mel spectrogram
        mel_spec = librosa.feature.melspectrogram(
            y=audio,
            sr=SAMPLE_RATE,
            n_fft=N_FFT,
            hop_length=HOP_LENGTH,
            n_mels=N_MELS,
            fmin=0,
            fmax=SAMPLE_RATE // 2
        )
        
        # Convert to log scale (dB)
        mel_spec_db = librosa.power_to_db(mel_spec, ref=np.max)
        
        # Normalize to [0, 1] range
        mel_spec_norm = (mel_spec_db - mel_spec_db.min()) / (mel_spec_db.max() - mel_spec_db.min() + 1e-8)
        
        # Debug: Check mel spec stats
        if hasattr(self, 'debug_count') and self.debug_count < 3:
            print(f"DEBUG - Mel spectrogram stats:")
            print(f"  Shape: {mel_spec_norm.shape}")
            print(f"  Min: {mel_spec_norm.min():.4f}, Max: {mel_spec_norm.max():.4f}")
            print(f"  Mean: {mel_spec_norm.mean():.4f}")
        
        # Ensure we have exactly N_FRAMES
        if mel_spec_norm.shape[1] < N_FRAMES:
            # Pad if too short
            padding = N_FRAMES - mel_spec_norm.shape[1]
            mel_spec_norm = np.pad(mel_spec_norm, ((0, 0), (0, padding)), mode='constant')
        elif mel_spec_norm.shape[1] > N_FRAMES:
            # Trim if too long
            mel_spec_norm = mel_spec_norm[:, :N_FRAMES]
        
        return mel_spec_norm
    
    def predict(self, mel_spec):
        """Run inference on mel spectrogram"""
        # Shape should be [1, 16, 96] for batch_size=1
        input_data = mel_spec.reshape(1, N_MELS, N_FRAMES).astype(np.float32)
        
        # Debug: Check input data
        if hasattr(self, 'debug_count') and self.debug_count < 3:
            print(f"DEBUG - Model input:")
            print(f"  Shape: {input_data.shape}")
            print(f"  Min: {input_data.min():.4f}, Max: {input_data.max():.4f}")
            print(f"  Mean: {input_data.mean():.4f}")
        
        # Run inference
        outputs = self.session.run(None, {self.input_name: input_data})
        
        # Debug: Check raw output
        if hasattr(self, 'debug_count') and self.debug_count < 3:
            print(f"DEBUG - Raw model output:")
            print(f"  Type: {type(outputs)}")
            print(f"  Length: {len(outputs)}")
            print(f"  Output[0] shape: {outputs[0].shape}")
            print(f"  Output[0]: {outputs[0]}")
            self.debug_count += 1
        
        # Get probability (assumes single output with probability)
        probability = outputs[0][0] if len(outputs[0].shape) > 1 else outputs[0]
        
        # Handle different output formats
        if isinstance(probability, np.ndarray):
            if probability.size > 1:
                # Binary classification - get positive class probability
                if hasattr(self, 'debug_count') and self.debug_count <= 3:
                    print(f"DEBUG - Binary output: {probability}")
                probability = probability[1] if len(probability) > 1 else probability[0]
            else:
                probability = probability.item()
        
        return float(probability)
    
    def run(self):
        """Main detection loop"""
        print(f"\nListening for wake word... (Threshold: {self.threshold})")
        print("Press Ctrl+C to stop\n")
        
        # Start audio stream
        with sd.InputStream(
            samplerate=SAMPLE_RATE,
            channels=1,
            blocksize=CHUNK_SIZE,
            callback=self.audio_callback
        ):
            try:
                while True:
                    # Get audio chunk from queue
                    audio_chunk = self.audio_queue.get()
                    
                    # Add to buffer
                    for sample in audio_chunk.flatten():
                        self.audio_buffer.append(sample)
                    
                    # Need enough audio for prediction
                    if len(self.audio_buffer) < self.required_samples:
                        continue
                    
                    # Get the last required_samples for processing
                    audio_window = np.array(list(self.audio_buffer)[-self.required_samples:])
                    
                    # Convert to mel spectrogram
                    mel_spec = self.audio_to_mel_spectrogram(audio_window)
                    
                    # Run prediction
                    probability = self.predict(mel_spec)
                    
                    # Check threshold
                    if probability > self.threshold:
                        print(f"🎤 WAKE WORD DETECTED! (confidence: {probability:.3f})")
                        # Clear buffer after detection to avoid repeated triggers
                        self.audio_buffer.clear()
                    else:
                        # Show activity indicator
                        sys.stdout.write(f"\rListening... {probability:.3f}")
                        sys.stdout.flush()
                        
            except KeyboardInterrupt:
                print("\n\nStopped listening.")

if __name__ == "__main__":
    # Check for model path argument
    if len(sys.argv) > 1:
        MODEL_PATH = sys.argv[1]
    
    # Optional threshold argument
    threshold = float(sys.argv[2]) if len(sys.argv) > 2 else THRESHOLD
    
    try:
        detector = WakeWordDetector(MODEL_PATH, threshold=threshold)
        detector.run()
    except FileNotFoundError:
        print(f"Error: Model file '{MODEL_PATH}' not found!")
        print("Usage: python script.py <model_path.onnx> [threshold]")
    except Exception as e:
        print(f"Error: {e}")
        import traceback
        traceback.print_exc()

im using above code to test but cannot detect wakeword
also probablity is constantly 0.00

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions