Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import pandas as pd | |
| import torch | |
| import torchaudio | |
| from transformers import Wav2Vec2ForSequenceClassification, Wav2Vec2Processor | |
| from sklearn.preprocessing import LabelEncoder | |
| import numpy as np | |
| # Set device | |
| device = torch.device("cuda" if torch.cuda.is_available() else "cpu") | |
| # Load the processor | |
| processor = Wav2Vec2Processor.from_pretrained("./fine_tuned_model") | |
| # Load the model | |
| model = Wav2Vec2ForSequenceClassification.from_pretrained("facebook/wav2vec2-base-960h", num_labels=len(pd.read_csv("dataset/train_wav.csv")["Common Name"].unique())) | |
| model.load_state_dict(torch.load("./fine_tuned_model/model.pt", map_location=device)) | |
| model.to(device) | |
| model.eval() | |
| # Load the label encoder | |
| label_encoder = LabelEncoder() | |
| 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): | |
| logits = torch.tensor(logits) # Convert numpy array to PyTorch tensor | |
| 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) | |
| # Ensure the audio is mono | |
| if waveform.size(0) > 1: | |
| waveform = torch.mean(waveform, dim=0, keepdim=True) | |
| # Ensure the audio is exactly 10 seconds | |
| if waveform.size(1) > fixed_length: | |
| waveform = waveform[:, :fixed_length] | |
| else: | |
| padding = fixed_length - waveform.size(1) | |
| waveform = torch.nn.functional.pad(waveform, (0, padding)) | |
| inputs = processor(waveform.squeeze(0), sampling_rate=16000, return_tensors="pt", padding=True).to(device) | |
| with torch.no_grad(): | |
| logits = model(inputs.input_values).logits | |
| return get_top_5_predictions(logits.cpu().numpy(), label_encoder) | |
| # Streamlit interface | |
| st.title("Bird Sound Classification") | |
| uploaded_file = st.file_uploader("Upload an audio file", type=["wav"]) | |
| if uploaded_file is not None: | |
| # Save the uploaded file temporarily | |
| file_path = f"temp/{uploaded_file.name}" | |
| os.makedirs(os.path.dirname(file_path), exist_ok=True) | |
| with open(file_path, "wb") as f: | |
| f.write(uploaded_file.getbuffer()) | |
| st.audio(file_path, format='audio/wav') | |
| if st.button("Predict"): | |
| with st.spinner("Classifying..."): | |
| 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}") |