Update app.py

app.py

@@ -3,51 +3,29 @@ import librosa
 import numpy as np
 import torch
 import logging
-from transformers import AutoModelForAudioClassification
+from transformers import Wav2Vec2ForSequenceClassification
 
 logging.basicConfig(level=logging.INFO)
 
 model_path = "./"
-model = AutoModelForAudioClassification.from_pretrained(model_path)
+model = Wav2Vec2ForSequenceClassification.from_pretrained(model_path)
 
-def augment_and_extract_features(audio_path, sr=16000, n_mfcc=40, n_fft=2048, hop_length=512, target_length=512):
-    y, sr = librosa.load(audio_path, sr=sr)
-    y_pitch_shifted = librosa.effects.pitch_shift(y, sr=sr, n_steps=4)
-    y_time_stretched = librosa.effects.time_stretch(y_pitch_shifted, rate=1.2)
-
-    mfcc = librosa.feature.mfcc(y=y_time_stretched, sr=sr, n_mfcc=n_mfcc, n_fft=n_fft, hop_length=hop_length)
-    chroma = librosa.feature.chroma_stft(y=y_time_stretched, sr=sr, n_fft=n_fft, hop_length=hop_length)
-    mel = librosa.feature.melspectrogram(y=y_time_stretched, sr=sr, n_fft=n_fft, hop_length=hop_length)
-    contrast = librosa.feature.spectral_contrast(y=y_time_stretched, sr=sr, n_fft=n_fft, hop_length=hop_length)
-    tonnetz = librosa.feature.tonnetz(y=librosa.effects.harmonic(y_time_stretched), sr=sr)
-
-    features = np.concatenate((mfcc, chroma, mel, contrast, tonnetz), axis=0)
-    features_normalized = (features - np.mean(features, axis=1, keepdims=True)) / np.std(features, axis=1, keepdims=True)
-
-    if features_normalized.shape[1] > target_length:
-        features_normalized = features_normalized[:, :target_length]
-    else:
-        padding = target_length - features_normalized.shape[1]
-        features_normalized = np.pad(features_normalized, ((0, 0), (0, padding)), 'constant')
-
-    features_tensor = torch.tensor(features_normalized).float().unsqueeze(0)  # Add batch dimension
-    return features_tensor
+def preprocess_audio(audio_path, target_sr=16000):
+    y, sr = librosa.load(audio_path, sr=target_sr)
+    y = librosa.resample(y, orig_sr=sr, target_sr=target_sr)
+    return y
 
 def predict_voice(audio_file_path):
     try:
-        features_tensor = augment_and_extract_features(audio_file_path)
-        # Correct the tensor shape to match expected input format for convolutional layers
-        if features_tensor.dim() > 4:  # Check if tensor has extra dimensions
-            features_tensor = features_tensor.squeeze()  # Remove unnecessary dimensions
-        if features_tensor.shape[-1] < model.config.num_labels:  # Ensure sufficient length for model input
-            padding_size = model.config.num_labels - features_tensor.shape[-1]
-            features_tensor = torch.nn.functional.pad(features_tensor, (0, padding_size), "constant", 0)
+        audio_data = preprocess_audio(audio_file_path)
+        inputs = model.processor(audio_data, sampling_rate=16000, return_tensors="pt", padding=True)
+
         with torch.no_grad():
-            outputs = model(features_tensor)
+            outputs = model(**inputs)
 
         logits = outputs.logits
-        predicted_index = logits.argmax()
-        label = model.config.id2label[predicted_index.item()]
+        predicted_index = logits.argmax(dim=1).item()
+        label = model.config.id2label[predicted_index]
         confidence = torch.softmax(logits, dim=1).max().item() * 100
 
         result = f"The voice is classified as '{label}' with a confidence of {confidence:.2f}%."