Update app.py

app.py

@@ -1,14 +1,13 @@
+import numpy as np
 import torch
-import torch.nn.functional as F
 import librosa
-import numpy as np
 import gradio as gr
 from transformers import AutoModelForAudioClassification
 import logging
 
 logging.basicConfig(level=logging.INFO)
 
-# Load your model here
+# Placeholder for loading your AST-compatible model
 model_path = "./"
 model = AutoModelForAudioClassification.from_pretrained(model_path)
 
@@ -17,26 +16,40 @@ def preprocess_audio(audio_path, sr=22050):
     audio, _ = librosa.effects.trim(audio)
     return audio, sr
 
+def extract_patches(S_DB, patch_size=16, patch_overlap=6):
+    stride = patch_size - patch_overlap
+    num_patches_x = (S_DB.shape[1] - patch_size) // stride + 1
+    num_patches_y = (S_DB.shape[0] - patch_size) // stride + 1
+    patches = []
+
+    for i in range(num_patches_y):
+        for j in range(num_patches_x):
+            start_i = i * stride
+            start_j = j * stride
+            patch = S_DB[start_i:start_i+patch_size, start_j:start_j+patch_size]
+            patches.append(patch.reshape(-1))
+    
+    return np.array(patches)
+
 def extract_features(audio, sr):
     S = librosa.feature.melspectrogram(y=audio, sr=sr, n_mels=128, hop_length=512, n_fft=2048)
     S_DB = librosa.power_to_db(S, ref=np.max)
 
-    # Reshape the spectrogram to a sequence of overlapping 16x16 patches
-    patches = librosa.util.frame(S_DB.flatten(), frame_length=16*16, hop_length=(16-6)*(16-6)).T
-    patches = patches.reshape(patches.shape[0], 16, 16)
-
-    # Linear projection layer equivalent (patch embedding layer)
-    patch_embeddings = patches.reshape(patches.shape[0], -1)
-    patch_embeddings = torch.tensor(patch_embeddings).float()
+    patches = extract_patches(S_DB)
+    patch_embeddings = torch.tensor(patches).float()
 
-    # Assuming positional embeddings and [CLS] token embedding are handled within the model
-    return patch_embeddings.unsqueeze(0)  # Add batch dimension for compatibility with model
+    # Assuming the model includes a patch embedding layer internally
+    return patch_embeddings.unsqueeze(0)  # Add batch dimension
 
 def predict_voice(audio_file_path):
     try:
         audio, sr = preprocess_audio(audio_file_path)
         features = extract_features(audio, sr)
         
+        # Flatten the patches to match the model's expected input shape
+        # Adjust this based on your AST model's input requirements
+        features = features.view(1, -1, 768)  # Reshape assuming the model expects (batch_size, seq_len, embedding_dim)
+        
         with torch.no_grad():
             outputs = model(features)
             logits = outputs.logits
@@ -57,7 +70,7 @@ iface = gr.Interface(
     inputs=gr.Audio(label="Upload Audio File", type="filepath"),
     outputs=gr.Text(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."
+    description="Detects whether a voice is real or AI-generated using an advanced AST model. Upload an audio file to see the results."
 )
 
 iface.launch()