Update app.py

app.py

@@ -2,6 +2,7 @@ import gradio as gr
 import librosa
 import numpy as np
 import torch
+import torch.nn.functional as F
 import logging
 from transformers import AutoModelForAudioClassification
 
@@ -12,7 +13,7 @@ logging.basicConfig(level=logging.INFO)
 local_model_path = "./"
 model = AutoModelForAudioClassification.from_pretrained(local_model_path)
 
-def custom_feature_extraction(audio_file_path, sr=16000, n_mels=128, n_fft=2048, hop_length=512):
+def custom_feature_extraction(audio_file_path, sr=16000, n_mels=128, n_fft=2048, hop_length=512, target_length=1024):
     """
     Custom feature extraction using Mel spectrogram, tailored for models trained on datasets like AudioSet.
     Args:
@@ -21,14 +22,25 @@ def custom_feature_extraction(audio_file_path, sr=16000, n_mels=128, n_fft=2048,
         n_mels: Number of Mel bands to generate.
         n_fft: Length of the FFT window.
         hop_length: Number of samples between successive frames.
+        target_length: Expected length of the Mel spectrogram in the time dimension.
     Returns:
         A tensor representation of the Mel spectrogram features.
     """
     waveform, sample_rate = librosa.load(audio_file_path, sr=sr)
     S = librosa.feature.melspectrogram(y=waveform, sr=sample_rate, n_mels=n_mels, n_fft=n_fft, hop_length=hop_length)
     S_DB = librosa.power_to_db(S, ref=np.max)
-    S_DB_tensor = torch.tensor(S_DB).float().unsqueeze(0)  # Add batch dimension
-    return S_DB_tensor
+    mel_tensor = torch.tensor(S_DB).float()
+
+    # Ensure the tensor matches the expected sequence length
+    current_length = mel_tensor.shape[1]
+    if current_length > target_length:
+        mel_tensor = mel_tensor[:, :target_length]  # Truncate if longer
+    elif current_length < target_length:
+        padding = target_length - current_length
+        mel_tensor = F.pad(mel_tensor, (0, padding), "constant", 0)  # Pad if shorter
+
+    mel_tensor = mel_tensor.unsqueeze(0)  # Add batch dimension for compatibility with model
+    return mel_tensor
 
 def predict_voice(audio_file_path):
     """
@@ -40,11 +52,8 @@ def predict_voice(audio_file_path):
     """
     try:
         features = custom_feature_extraction(audio_file_path)
-        
         with torch.no_grad():
-            # Corrected: Directly pass the features tensor to the model
             outputs = model(features)
-        
         logits = outputs.logits
         predicted_index = logits.argmax()
         label = model.config.id2label[predicted_index.item()]
@@ -68,4 +77,4 @@ iface = gr.Interface(
 )
 
 # Launching the interface
-iface.launch()
+iface.launch()