Update app.py

app.py

@@ -6,8 +6,11 @@ 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
+# Define the BirdCallRNN model class
 class BirdCallRNN(nn.Module):
     def __init__(self, resnet, num_features, num_classes):
         super(BirdCallRNN, self).__init__()
@@ -21,77 +24,87 @@ 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, :])  # Note: We’ll use this for single-segment sequences
+        output = self.fc(rnn_out[:, -1, :])
         return output
 
-# 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
+# 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
 
-# Load class mapping globally
+# Load class mapping
 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)  # Should be 114
+num_classes = len(class_names)
 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 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
+# 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
 
-# Launch Gradio interface
+# 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")
+    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'>
+    """
 )
 interface.launch()