Update app.py

app.py

@@ -10,51 +10,31 @@ logging.basicConfig(level=logging.INFO)
 model_path = "./"
 model = AutoModelForAudioClassification.from_pretrained(model_path)
 
-def preprocess_audio(audio_path, sr=22050):
+def preprocess_audio(audio_path, sr=16000):
+    # Load the audio file. Note: Adjusting the sample rate (sr) to match the model's expected input
     audio, sr = librosa.load(audio_path, sr=sr)
+    # Trim silence from the beginning and the end
     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_time = (S_DB.shape[1] - patch_overlap) // stride
-    num_patches_freq = (S_DB.shape[0] - patch_overlap) // stride
-    
-    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):
+def extract_features(audio, sr=16000):
+    # Compute the Mel spectrogram
     S = librosa.feature.melspectrogram(y=audio, sr=sr, n_mels=128, hop_length=512, n_fft=2048)
+    # Convert to dB scale
     S_DB = librosa.power_to_db(S, ref=np.max)
-    patches = extract_patches(S_DB)
-    
-    # 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
+    # Normally, further feature extraction steps would be here. For this model, we will directly use S_DB.
+    return S_DB
 
 def predict_voice(audio_file_path):
     try:
         audio, sr = preprocess_audio(audio_file_path)
-        features = extract_features(audio, sr)
+        S_DB = extract_features(audio, sr)
         
-        # 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)
+        # Convert S_DB to tensor and add required batch dimension
+        S_DB_tensor = torch.tensor(S_DB).unsqueeze(0)
         
         with torch.no_grad():
-            outputs = model(features)
+            outputs = model(S_DB_tensor)
             logits = outputs.logits
             predicted_index = logits.argmax()
             label = model.config.id2label[predicted_index.item()]