Predict.py

"""
Inference script for speech recognition using a trained SpeechRecognitionModel.

This script:
- Loads a trained CTC-based speech recognition model
- Preprocesses an input audio file into log-mel spectrograms
- Runs inference on the audio
- Decodes predictions using greedy CTC decoding
"""
import torch
import torchaudio
import torch.nn.functional as F
from VoiceRecognition import SpeechRecognitionModel  # adjust if yours is named differently

# Load your index-to-character map
idx_to_char = {
    0: '<pad>',
    1: '<unk>',
    2: 'a',
    3: 'b',
    4: 'c',
    5: 'd',
    6: 'e',
    7: 'f',
    8: 'g',
    9: 'h',
    10: 'i',
    11: 'j',
    12: 'k',
    13: 'l',
    14: 'm',
    15: 'n',
    16: 'o',
    17: 'p',
    18: 'q',
    19: 'r',
    20: 's',
    21: 't',
    22: 'u',
    23: 'v',
    24: 'w',
    25: 'x',
    26: 'y',
    27: 'z',
    28: ' ',
    29: "'"}  # fill in your map

def greedy_decoder(output, blank=0):
    """
        Perform greedy CTC decoding on model output.

        Args:
            output (Tensor): Model output logits of shape (B, T, C).
            blank (int): Index of the CTC blank token.

        Returns:
            list[str]: Decoded strings for each batch element.
    """
    arg_maxes = torch.argmax(output, dim=2)
    decodes = []
    for i, args in enumerate(arg_maxes):
        decode = []
        prev = -1
        for j in range(len(args)):
            if args[j] != prev:
                if args[j] != blank:
                    decode.append(idx_to_char[args[j].item()])
                prev = args[j]
        decodes.append(''.join(decode))
    return decodes

def load_model(checkpoint_path, device):
    """
       Load a trained speech recognition model from a checkpoint.

       Args:
           checkpoint_path (str): Path to the model checkpoint (.pth).
           device (torch.device): Device to load the model onto.

       Returns:
           SpeechRecognitionModel: Loaded model in evaluation mode.
    """
    checkpoint = torch.load(checkpoint_path, map_location=device)

    vocab_size = 30
    model = SpeechRecognitionModel(output_size=vocab_size)# ← your model class
    model.load_state_dict(checkpoint["model_state_dict"])
    model.to(device)
    model.eval()
    return model

def preprocess_audio(audio_path, transform, amplitude_to_db, device):
    """
        Load and preprocess an audio file into a log-mel spectrogram.

        Steps:
        - Load waveform
        - Convert to mono if needed
        - Resample to 16 kHz
        - Convert to mel spectrogram
        - Convert amplitude to decibels

        Args:
            audio_path (str): Path to the audio file (.wav).
            transform (MelSpectrogram): Mel spectrogram transform.
            amplitude_to_db (AmplitudeToDB): Log-amplitude transform.
            device (torch.device): Device to move tensor to.

        Returns:
            Tensor: Log-mel spectrogram of shape (1, 1, n_mels, time).
    """
    waveform, sample_rate = torchaudio.load(audio_path)
    if waveform.shape[0] > 1:
        waveform = torch.mean(waveform, dim=0, keepdim=True)
    if sample_rate != 16000:
        resample = torchaudio.transforms.Resample(orig_freq=sample_rate, new_freq=16000)
        waveform = resample(waveform)

    mel = transform(waveform)           # [1, 64, time]
    mel_db = amplitude_to_db(mel)       # log-mel
    return mel_db.unsqueeze(0).to(device)  # [1, 1, 64, time]

def predict(audio_path, model, mel_transform, amplitude_to_db, device):
    """
        Run speech recognition on an audio file.

        Args:
            audio_path (str): Path to the audio file.
            model (SpeechRecognitionModel): Trained model.
            mel_transform (MelSpectrogram): Mel spectrogram transform.
            amplitude_to_db (AmplitudeToDB): Log-amplitude transform.
            device (torch.device): Device to run inference on.

        Returns:
            str: Transcribed text output.
    """
    features = preprocess_audio(audio_path, mel_transform, amplitude_to_db, device)
    with torch.no_grad():
        output = model(features)  # [T, B, C]
        output = F.log_softmax(output, dim=2)
    decoded = greedy_decoder(output.cpu())
    return decoded[0]

if __name__ == "__main__":
    """
        Command-line entry point for speech recognition inference.

        Usage:
            python infer.py path/to/audio.wav
    """
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("audio", help="Path to audio file (.wav)")
    args = parser.parse_args()

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model = load_model("checkpointGPU10other.pth", device)

    mel_transform = torchaudio.transforms.MelSpectrogram(
        sample_rate=16000,
        n_fft=1024,
        hop_length=512,
        n_mels=64
    )
    amplitude_to_db = torchaudio.transforms.AmplitudeToDB()

    transcription = predict(args.audio, model, mel_transform,amplitude_to_db, device)
    print("Prediction:", transcription.replace("<pad>", " ").replace("<unk>", " "))