Update app.py

app.py

@@ -1,64 +1,97 @@
+import gradio as gr
 import torch
-import torchaudio
-import torchaudio.transforms as transforms
-import torch.nn.functional as F
+from torch import nn
+import cv2
+import numpy as np
 import json
-from flask import Flask, request, jsonify
 from torchvision import models
 import librosa
-import numpy as np
 
-app = Flask(__name__)
+# Define the BirdCallRNN model
+class BirdCallRNN(nn.Module):
+    def __init__(self, resnet, num_features, num_classes):
+        super(BirdCallRNN, self).__init__()
+        self.resnet = resnet
+        self.rnn = nn.LSTM(input_size=num_features, hidden_size=256, num_layers=2, batch_first=True, bidirectional=True)
+        self.fc = nn.Linear(512, num_classes)
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    def forward(self, x):
+        batch, seq_len, C, H, W = x.size()
+        x = x.view(batch * seq_len, C, H, W)
+        features = self.resnet(x)
+        features = features.view(batch, seq_len, -1)
+        rnn_out, _ = self.rnn(features)
+        output = self.fc(rnn_out[:, -1, :])  # Note: We’ll use this for single-segment sequences
+        return output
 
-# Load ResNet50 feature extractor
-resnet = models.resnet50(pretrained=False)
-resnet.fc = torch.nn.Identity()  # Remove classification layer
-resnet = resnet.to(device)
-resnet.eval()
+# Function to convert MP3 to mel spectrogram (unchanged)
+def mp3_to_mel_spectrogram(mp3_file, target_shape=(128, 500), resize_shape=(224, 224)):
+    y, sr = librosa.load(mp3_file, sr=None)
+    S = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=128, fmax=8000)
+    log_S = librosa.power_to_db(S, ref=np.max)
+    current_time_steps = log_S.shape[1]
+    target_time_steps = target_shape[1]
+    if current_time_steps < target_time_steps:
+        pad_width = target_time_steps - current_time_steps
+        log_S_resized = np.pad(log_S, ((0, 0), (0, pad_width)), mode='constant')
+    elif current_time_steps > target_time_steps:
+        log_S_resized = log_S[:, :target_time_steps]
+    else:
+        log_S_resized = log_S
+    log_S_resized = cv2.resize(log_S_resized, resize_shape, interpolation=cv2.INTER_CUBIC)
+    return log_S_resized
 
-# Load RNN model for classification
-rnn_model = torch.nn.LSTM(input_size=2048, hidden_size=512, num_layers=2, batch_first=True, bidirectional=True)
-rnn_fc = torch.nn.Linear(1024, 114)  # 114 classes
-rnn_model = rnn_model.to(device)
-rnn_fc = rnn_fc.to(device)
-rnn_model.eval()
-rnn_fc.eval()
+# Load class mapping globally
+with open('class_mapping.json', 'r') as f:
+    class_names = json.load(f)
 
-# Load class mapping
-with open("class_mapping.json", "r") as f:
-    class_mapping = json.load(f)
-
-def preprocess_audio(file_path):
-    y, sr = librosa.load(file_path, sr=22050)  # Load at 22.05kHz
-    mel_spectrogram = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=128, fmax=8000)
-    mel_spectrogram = librosa.power_to_db(mel_spectrogram, ref=np.max)
-    mel_spectrogram = torch.tensor(mel_spectrogram).unsqueeze(0).unsqueeze(0)  # Add batch & channel dims
-    return mel_spectrogram.to(device)
+# Revised inference function to predict per segment
+def infer_birdcall(model, mp3_file, segment_length=500, device="cuda"):
+    model.eval()
+    # Load audio and compute mel spectrogram
+    y, sr = librosa.load(mp3_file, sr=None)
+    S = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=128, fmax=8000)
+    log_S = librosa.power_to_db(S, ref=np.max)
+    # Segment the spectrogram
+    num_segments = log_S.shape[1] // segment_length
+    if num_segments == 0:
+        segments = [log_S]
+    else:
+        segments = [log_S[:, i * segment_length:(i + 1) * segment_length] for i in range(num_segments)]
+    
+    predictions = []
+    # Process each segment individually
+    for seg in segments:
+        seg_resized = cv2.resize(seg, (224, 224), interpolation=cv2.INTER_CUBIC)
+        seg_rgb = np.repeat(seg_resized[:, :, np.newaxis], 3, axis=-1)
+        # Create a tensor with batch size 1 and sequence length 1
+        seg_tensor = torch.from_numpy(seg_rgb).permute(2, 0, 1).float().unsqueeze(0).unsqueeze(0).to(device)  # Shape: (1, 1, 3, 224, 224)
+        output = model(seg_tensor)
+        pred = torch.max(output, dim=1)[1].cpu().numpy()[0]
+        predicted_bird = class_names[str(pred)]  # Convert pred to string to match JSON keys
+        predictions.append(predicted_bird)
+    return predictions
 
-def predict_birdcall(audio_tensor):
-    with torch.no_grad():
-        features = resnet(audio_tensor.repeat(1, 3, 1, 1))  # Convert grayscale to 3-channel
-        features = features.unsqueeze(1)  # Add sequence dimension
-        rnn_out, _ = rnn_model(features)
-        logits = rnn_fc(rnn_out[:, -1, :])  # Take last time step
-        pred = torch.argmax(logits, dim=1).item()
-        return class_mapping.get(str(pred), "Unknown")
+# Initialize the model
+resnet = models.resnet50(weights='IMAGENET1K_V2')
+num_features = resnet.fc.in_features
+resnet.fc = nn.Identity()
+num_classes = len(class_names)  # Should be 114
+model = BirdCallRNN(resnet, num_features, num_classes)
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model.to(device)
+model.load_state_dict(torch.load('model_weights.pth', map_location=device))
+model.eval()
 
-@app.route("/predict", methods=["POST"])
-def predict():
-    if "file" not in request.files:
-        return jsonify({"error": "No file uploaded"}), 400
-    file = request.files["file"]
-    
-    file_path = "temp_audio.wav"
-    file.save(file_path)
-    
-    audio_tensor = preprocess_audio(file_path)
-    prediction = predict_birdcall(audio_tensor)
-    
-    return jsonify({"prediction": prediction})
+# Prediction function for Gradio
+def predict_bird(file_path):
+    predictions = infer_birdcall(model, file_path, segment_length=500, device=str(device))
+    return ", ".join(predictions)  # Join predictions into a single string
 
-if __name__ == "__main__":
-    app.run(host="0.0.0.0", port=5000, debug=True)
+# Launch Gradio interface
+interface = gr.Interface(
+    fn=predict_bird,
+    inputs=gr.File(label="Upload MP3 file", file_types=['.mp3']),
+    outputs=gr.Textbox(label="Predicted Bird Species")
+)
+interface.launch()