Update app.py

app.py

@@ -1,49 +1,54 @@
+import gradio as gr
+import librosa
 import numpy as np
 import torch
-import librosa
-import gradio as gr
-from transformers import AutoModelForAudioClassification, Wav2Vec2Processor
-import logging
+import torch.nn.functional as F
+from transformers import AutoModelForAudioClassification
+import random
 
-logging.basicConfig(level=logging.INFO)
+model = AutoModelForAudioClassification.from_pretrained("./")
 
-model_path = "./"
-model = AutoModelForAudioClassification.from_pretrained(model_path)
-processor = Wav2Vec2Processor.from_pretrained(model_path)
+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)
+    pitches, _ = librosa.piptrack(y=waveform, sr=sample_rate, n_fft=n_fft, hop_length=hop_length)
+    spectral_centroids = librosa.feature.spectral_centroid(y=waveform, sr=sample_rate, n_fft=n_fft, hop_length=hop_length)
+    features = np.concatenate([S_DB, pitches, 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 preprocess_audio(audio_path, sr=16000):
-    audio, _ = librosa.load(audio_path, sr=sr)
-    audio, _ = librosa.effects.trim(audio)
-    return audio
-
-def extract_features(audio, sr=16000):
-    inputs = processor(audio, sampling_rate=sr, return_tensors="pt", padding=True)
-    return inputs
+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)))
+    return np.roll(waveform, shift)
 
 def predict_voice(audio_file_path):
     try:
-        audio = preprocess_audio(audio_file_path)
-        features = extract_features(audio)
-        
+        waveform, sample_rate = librosa.load(audio_file_path, sr=None)
+        augmented_waveform = apply_time_shift(waveform)
+        original_features = custom_feature_extraction(audio_file_path, sr=sample_rate)
+        augmented_features = custom_feature_extraction(augmented_waveform, sr=sample_rate)
         with torch.no_grad():
-            outputs = model(**features)
-            logits = outputs.logits
-            predicted_index = logits.argmax(dim=-1)
-            label = processor.decode(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}%."
-        logging.info("Prediction successful.")
+            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()]
+        confidence = torch.softmax(logits, dim=1).max().item() * 100
+        return f"The voice is classified as '{label}' with a confidence of {confidence:.2f}%."
     except Exception as e:
-        result = f"Error during processing: {e}"
-        logging.error(result)
-
-    return result
+        return f"Error during processing: {e}"
 
 iface = gr.Interface(
     fn=predict_voice,
     inputs=gr.Audio(label="Upload Audio File", type="filepath"),
-    outputs=gr.Text(label="Prediction"),
+    outputs=gr.Textbox(label="Prediction"),
     title="Voice Authenticity Detection",
-    description="This system uses advanced audio processing to detect whether a voice is real or AI-generated. Upload an audio file to see the results."
-).launch()
+    description="Detects whether a voice is real or AI-generated. Upload an audio file to see the results."
+)
+
+iface.launch()