Update app.py

app.py

@@ -6,11 +6,8 @@ import numpy as np
 import json
 from torchvision import models
 import librosa
-import matplotlib.pyplot as plt
-from io import BytesIO
-import PIL.Image
 
-# Define the BirdCallRNN model class
+# Define the BirdCallRNN model
 class BirdCallRNN(nn.Module):
     def __init__(self, resnet, num_features, num_classes):
         super(BirdCallRNN, self).__init__()
@@ -24,87 +21,77 @@ class BirdCallRNN(nn.Module):
         features = self.resnet(x)
         features = features.view(batch, seq_len, -1)
         rnn_out, _ = self.rnn(features)
-        output = self.fc(rnn_out[:, -1, :])
+        output = self.fc(rnn_out[:, -1, :])  # Note: We’ll use this for single-segment sequences
         return output
 
-# Function to plot mel spectrogram
-def plot_spectrogram(log_S, sr):
-    fig, ax = plt.subplots(figsize=(10, 4))
-    img = librosa.display.specshow(log_S, sr=sr, x_axis='time', y_axis='mel', ax=ax)
-    fig.colorbar(img, ax=ax, format='%+2.0f dB')
-    ax.set_title('Mel Spectrogram')
-    buf = BytesIO()
-    plt.savefig(buf, format='png')
-    buf.seek(0)
-    img = PIL.Image.open(buf)
-    plt.close(fig)
-    return img
+# 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 class mapping
+# Load class mapping globally
 with open('class_mapping.json', 'r') as f:
     class_names = json.load(f)
 
+# 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
+
 # Initialize the model
 resnet = models.resnet50(weights='IMAGENET1K_V2')
 num_features = resnet.fc.in_features
 resnet.fc = nn.Identity()
-num_classes = len(class_names)
+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()
 
-# Prediction function
-def predict_bird(audio):
-    # Load audio file
-    y, sr = librosa.load(audio, 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)
-    
-    # Generate spectrogram image
-    spectrogram_img = plot_spectrogram(log_S, sr)
-    
-    # Segment audio and predict
-    predictions = []
-    segment_length = 500
-    num_segments = log_S.shape[1] // segment_length
-    segments = [log_S] if num_segments == 0 else [log_S[:, i * segment_length:(i + 1) * segment_length] for i in range(num_segments)]
-    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)
-        seg_tensor = torch.from_numpy(seg_rgb).permute(2, 0, 1).float().unsqueeze(0).unsqueeze(0).to(device)
-        output = model(seg_tensor)
-        probs = torch.softmax(output, dim=1)
-        confidence, pred = torch.max(probs, dim=1)
-        pred = pred.cpu().numpy()[0]
-        confidence = confidence.cpu().numpy()[0]
-        predicted_bird = class_names[str(pred)]
-        predictions.append((predicted_bird, confidence))
-    
-    # Format predictions as HTML
-    predictions_html = "<ol>"
-    for i, (bird, conf) in enumerate(predictions, 1):
-        predictions_html += f"<li>{bird} (Confidence: {conf*100:.1f}%)</li>"
-    predictions_html += "</ol>"
-    
-    return spectrogram_img, predictions_html
+# 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
 
-# Gradio interface
+# Launch Gradio interface
 interface = gr.Interface(
     fn=predict_bird,
-    inputs=gr.Audio(label="Upload MP3 file", type="filepath"),
-    outputs=[
-        gr.Image(label="Mel Spectrogram"),
-        gr.HTML(label="Predicted Bird Species")
-    ],
-    description="""
-    <h3>Bird Species</h3>
-    <img src='1.jpeg' width='300'>
-    <h3>Bird Description</h3>
-    <img src='2.jpeg' width='300'>
-    <h3>Bird Origins</h3>
-    <img src='3.jpeg' width='300'>
-    """
+    inputs=gr.File(label="Upload MP3 file", file_types=['.mp3']),
+    outputs=gr.Textbox(label="Predicted Bird Species")
 )
 interface.launch()