Update app.py

app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import pandas as pd
 import torch
 import torchaudio
-from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2Processor
+from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2Processor, safetensors
 from sklearn.preprocessing import LabelEncoder
 import numpy as np
 
@@ -10,7 +10,7 @@ import numpy as np
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 # Load the fine-tuned model and processor
-model = Wav2Vec2ForSequenceClassification.from_pretrained("./fine_tuned_model").to(device)
+model = Wav2Vec2ForSequenceClassification.from_pretrained("./fine_tuned_model", from_safetensors=True).to(device)
 processor = Wav2Vec2Processor.from_pretrained("./fine_tuned_model")
 
 # Load the label encoder
@@ -20,6 +20,15 @@ label_encoder.fit(pd.read_csv("dataset/train_wav.csv")["Common Name"])
 # Fixed audio length (e.g., 10 seconds)
 fixed_length = 10 * 16000  # 10 seconds * 16000 Hz
 
+# Function to get top 5 predictions with probabilities
+def get_top_5_predictions(logits, label_encoder):
+    probabilities = torch.nn.functional.softmax(logits, dim=-1).cpu().numpy()
+    top5_idx = np.argsort(probabilities, axis=-1)[:, -5:][:, ::-1]  # Top 5 indices
+    top5_probs = np.take_along_axis(probabilities, top5_idx, axis=-1)
+    top5_labels = label_encoder.inverse_transform(top5_idx[0])
+    
+    return list(zip(top5_labels, top5_probs[0]))
+
 # Prediction function
 def predict(file_path):
     waveform, sample_rate = torchaudio.load(file_path)
@@ -40,12 +49,7 @@ def predict(file_path):
     with torch.no_grad():
         logits = model(inputs.input_values).logits
     
-    probabilities = torch.nn.functional.softmax(logits, dim=-1).cpu().numpy()[0]
-    top5_idx = np.argsort(probabilities)[-5:][::-1]
-    top5_probs = probabilities[top5_idx]
-    top5_labels = label_encoder.inverse_transform(top5_idx)
-    
-    return list(zip(top5_labels, top5_probs))
+    return get_top_5_predictions(logits, label_encoder)
 
 # Streamlit interface
 st.title("Bird Sound Classification")
@@ -66,4 +70,4 @@ if uploaded_file is not None:
             top5_predictions = predict(file_path)
             st.success("Top 5 Predicted Bird Species with Probabilities:")
             for label, prob in top5_predictions:
-                st.write(f"{label}: {prob:.4f}")
+                st.write(f"{label}: {prob:.4f}")