Update app.py

app.py

@@ -7,7 +7,6 @@ import logging
 
 logging.basicConfig(level=logging.INFO)
 
-# Placeholder for loading your AST-compatible model
 model_path = "./"
 model = AutoModelForAudioClassification.from_pretrained(model_path)
 
@@ -18,37 +17,41 @@ def preprocess_audio(audio_path, sr=22050):
 
 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))
+    num_patches_time = (S_DB.shape[1] - patch_overlap) // stride
+    num_patches_freq = (S_DB.shape[0] - patch_overlap) // stride
     
-    return np.array(patches)
+    patches = []
+    for i in range(0, num_patches_freq * stride, stride):
+        for j in range(0, num_patches_time * stride, stride):
+            patch = S_DB[i:i+patch_size, j:j+patch_size]
+            if patch.shape == (patch_size, patch_size):
+                patches.append(patch.reshape(-1))
+    return np.stack(patches) if patches else np.empty((0, patch_size*patch_size))
 
 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)
-
     patches = extract_patches(S_DB)
-    patch_embeddings = torch.tensor(patches).float()
-
-    # Assuming the model includes a patch embedding layer internally
-    return patch_embeddings.unsqueeze(0)  # Add batch dimension
+    
+    # Assuming each patch is flattened to a vector of size 256 (16*16) and then projected to 768 dimensions
+    # Here we simulate this projection by creating a dummy tensor, in practice, this should be done by a learned linear layer
+    patches_tensor = torch.tensor(patches).float()
+    # Simulate linear projection (e.g., via a fully connected layer) to match the embedding size
+    if patches_tensor.nelement() == 0:  # Handle case of no patches
+        patch_embeddings_tensor = torch.empty(0, 768)
+    else:
+        patch_embeddings_tensor = patches_tensor  # This is a placeholder, replace with actual projection
+    
+    return patch_embeddings_tensor.unsqueeze(0)  # Add batch dimension for compatibility with model
 
 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)
+        # Adjust the features size to match the model input, if necessary
+        # Example: Reshape or pad the features tensor
+        # features = adjust_features_shape(features, expected_shape)
         
         with torch.no_grad():
             outputs = model(features)
@@ -70,7 +73,7 @@ iface = gr.Interface(
     inputs=gr.Audio(label="Upload Audio File", type="filepath"),
     outputs=gr.Text(label="Prediction"),
     title="Voice Authenticity Detection",
-    description="Detects whether a voice is real or AI-generated using an advanced AST model. Upload an audio file to see the results."
+    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."
 )
 
 iface.launch()