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,19 +13,21 @@ 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):
     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
+    
+    # Resizing the tensor to match the model's expected input size
+    S_DB_tensor_resized = F.interpolate(S_DB_tensor, size=(n_mels, target_length), mode='nearest')
+    return S_DB_tensor_resized
 
 def predict_voice(audio_file_path):
     try:
         features = custom_feature_extraction(audio_file_path)
         
         with torch.no_grad():
-            # Directly pass the features tensor to the model
             outputs = model(features)
         
         logits = outputs.logits