Update app.py

app.py

@@ -3,27 +3,17 @@ import librosa
 import numpy as np
 import torch
 import torch.nn.functional as F
-from transformers import AutoModelForAudioClassification
+from transformers import AutoModelForAudioClassification, ASTFeatureExtractor
 import random
 
+# Model and feature extractor loading from the specified local path
 model = AutoModelForAudioClassification.from_pretrained("./")
+feature_extractor = ASTFeatureExtractor.from_pretrained("./")
 
-def custom_feature_extraction(audio, sr=16000, n_mels=128, n_fft=2048, hop_length=512, target_length=1024):
-    S = librosa.feature.melspectrogram(y=audio, sr=sr, n_mels=n_mels, n_fft=n_fft, hop_length=hop_length)
-    S_DB = librosa.power_to_db(S, ref=np.max)
-    pitches, _ = librosa.piptrack(y=audio, sr=sr, n_fft=n_fft, hop_length=hop_length)
-    spectral_centroids = librosa.feature.spectral_centroid(y=audio, sr=sr, n_fft=n_fft, hop_length=hop_length)
-    # Correct the dimensionality issue
-    pitches_max = np.max(pitches, axis=0, keepdims=True)
-    spectral_centroids = spectral_centroids.reshape(1, -1)
-    # Ensure the concatenation axis has matching dimensions
-    features = np.concatenate([S_DB, pitches_max, spectral_centroids], axis=0)
-    features_tensor = torch.tensor(features).float()
-    if features_tensor.shape[1] > target_length:
-        features_tensor = features_tensor[:, :target_length]
-    else:
-        features_tensor = F.pad(features_tensor, (0, target_length - features_tensor.shape[1]), 'constant', 0)
-    return features_tensor.unsqueeze(0)
+def custom_feature_extraction(audio, sr=16000, n_mels=128, target_length=1024):
+    # Using the loaded feature extractor
+    features = feature_extractor(audio, sampling_rate=sr, return_tensors="pt", padding="max_length", max_length=target_length)
+    return features.input_values
 
 def apply_time_shift(waveform, max_shift_fraction=0.1):
     shift = random.randint(-int(max_shift_fraction * len(waveform)), int(max_shift_fraction * len(waveform)))
@@ -31,13 +21,16 @@ def apply_time_shift(waveform, max_shift_fraction=0.1):
 
 def predict_voice(audio_file_path):
     try:
-        waveform, sample_rate = librosa.load(audio_file_path, sr=None)
+        waveform, sample_rate = librosa.load(audio_file_path, sr=feature_extractor.sampling_rate, mono=True)
         augmented_waveform = apply_time_shift(waveform)
+        
         original_features = custom_feature_extraction(waveform, sr=sample_rate)
         augmented_features = custom_feature_extraction(augmented_waveform, sr=sample_rate)
+        
         with torch.no_grad():
             outputs_original = model(original_features)
             outputs_augmented = model(augmented_features)
+        
         logits = (outputs_original.logits + outputs_augmented.logits) / 2
         predicted_index = logits.argmax()
         label = model.config.id2label[predicted_index.item()]
@@ -55,4 +48,3 @@ iface = gr.Interface(
 )
 
 iface.launch()
-