Update app.py

app.py

@@ -2,65 +2,38 @@ import gradio as gr
 import librosa
 import numpy as np
 import torch
-import torch.nn.functional as F
 import logging
 from transformers import AutoModelForAudioClassification
-import random
 
-# Configure logging
 logging.basicConfig(level=logging.INFO)
 
-# Load the model
 model_path = "./"
 model = AutoModelForAudioClassification.from_pretrained(model_path)
 
-def custom_feature_extraction(waveform, sr, n_mels=128, n_fft=2048, hop_length=512, target_length=1024):
-    # Generate Mel spectrogram
+def preprocess_audio(audio_file_path, sr=16000):
+    waveform, _ = librosa.load(audio_file_path, sr=sr)
+    waveform = librosa.effects.trim(waveform)[0]  # Trim silence
+    return waveform
+
+def extract_features(waveform, sr=16000, n_mels=128, n_fft=2048, hop_length=512):
     S = librosa.feature.melspectrogram(y=waveform, sr=sr, n_mels=n_mels, n_fft=n_fft, hop_length=hop_length)
     S_DB = librosa.power_to_db(S, ref=np.max)
-    
-    # Pitch feature (using piptrack to estimate pitches and then averaging)
-    pitches, _ = librosa.piptrack(y=waveform, sr=sr, n_fft=n_fft, hop_length=hop_length)
-    pitch_mean = np.mean(pitches, axis=0, keepdims=True)
-
-    # Spectral centroid
-    spectral_centroids = librosa.feature.spectral_centroid(y=waveform, sr=sr, n_fft=n_fft, hop_length=hop_length)
-
-    # Concatenate features and normalize
-    features = np.concatenate([S_DB, pitch_mean, spectral_centroids], axis=0)
-    features_tensor = torch.tensor(features).float()
-
-    # Adjust the tensor's length
-    if features_tensor.shape[1] > target_length:
-        features_tensor = features_tensor[:, :target_length]
-    elif features_tensor.shape[1] < target_length:
-        padding = target_length - features_tensor.shape[1]
-        features_tensor = F.pad(features_tensor, (0, padding), 'constant', 0)
-
-    return features_tensor.unsqueeze(0)
+    return torch.tensor(S_DB).float().unsqueeze(0)  # Add batch dimension
 
-def apply_time_shift(waveform, max_shift_fraction=0.1):
-    shift_amount = int(max_shift_fraction * len(waveform))
-    shift = random.randint(-shift_amount, shift_amount)
-    return np.roll(waveform, shift)
-
-def predict_voice(audio_file):
+def predict_voice(audio_file_path):
     try:
-        waveform, sample_rate = librosa.load(audio_file, sr=None)
-        augmented_waveform = apply_time_shift(waveform)
-
-        original_features = custom_feature_extraction(waveform, sample_rate)
-        augmented_features = custom_feature_extraction(augmented_waveform, sample_rate)
-
+        waveform = preprocess_audio(audio_file_path)
+        features = extract_features(waveform)
+        
         with torch.no_grad():
-            outputs_original = model(original_features)
-            outputs_augmented = model(augmented_features)
-            logits = (outputs_original.logits + outputs_augmented.logits) / 2
+            outputs = model(features)
+            logits = outputs.logits
             predicted_index = logits.argmax()
             label = model.config.id2label[predicted_index.item()]
             confidence = torch.softmax(logits, dim=1).max().item() * 100
-
+        
         result = f"The voice is classified as '{label}' with a confidence of {confidence:.2f}%."
+        logging.info("Prediction successful.")
     except Exception as e:
         result = f"Error during processing: {e}"
         logging.error(result)
@@ -69,8 +42,8 @@ def predict_voice(audio_file):
 
 iface = gr.Interface(
     fn=predict_voice,
-    inputs=gr.Audio(label="Upload Audio File", type="filepath"),
-    outputs=gr.Textbox(label="Prediction"),
+    inputs=gr.Audio(label="Upload Audio File", type="file"),
+    outputs=gr.Text(label="Prediction"),
     title="Voice Authenticity Detection",
     description="Detects whether a voice is real or AI-generated. Upload an audio file to see the results."
 )