Update app.py

app.py

@@ -21,61 +21,63 @@ 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 convert MP3 to mel spectrogram
+# 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)
-
-    # Ensure the correct time step size
     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')
-    else:
+    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)
 
-# Inference function for bird call identification
+# 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
-    segments = [log_S[:, i * segment_length:(i + 1) * segment_length] for i in range(num_segments)] if num_segments > 0 else [log_S]
-
+    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)
-        seg_tensor = torch.from_numpy(seg_rgb).permute(2, 0, 1).float().unsqueeze(0).unsqueeze(0).to(device)
+        # 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)]
+        predicted_bird = class_names[str(pred)]  # Convert pred to string to match JSON keys
         predictions.append(predicted_bird)
+    return predictions
 
-    return "\n".join([f"{i+1}. {pred}" for i, pred in enumerate(predictions)])
-
-# Initialize model
+# 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))
@@ -83,29 +85,37 @@ model.eval()
 
 # Prediction function for Gradio
 def predict_bird(file_path):
-    return infer_birdcall(model, file_path, segment_length=500, device=str(device))
-
-# Gradio interface
-with gr.Blocks() as interface:
-    gr.Markdown("### Bird Call Identification")
-
-    # File upload and audio preview
-    with gr.Row():
-        file_input = gr.File(label="Upload MP3 file", file_types=['.mp3'])
-        audio_player = gr.Audio(label="Uploaded MP3 File")
-
-    # Update audio player after upload
-    file_input.change(lambda file: file.name if file else None, inputs=file_input, outputs=audio_player)
-
-    # Prediction UI
-    with gr.Row():
-        submit_button = gr.Button("Submit")
-
-    prediction_output = gr.Textbox(label="Predicted Bird Species")
-    bird_species_image = gr.Image(label="Bird Species")
-    bird_description_image = gr.Image(label="Bird Description")
-    bird_origins_image = gr.Image(label="Bird Origins")
-
-    submit_button.click(predict_bird, inputs=file_input, outputs=prediction_output)
-
-interface.launch()
+    predictions = infer_birdcall(model, file_path, segment_length=500, device=str(device))
+    # Format predictions as a numbered list
+    formatted_predictions = "\n".join([f"{i+1}. {pred}" for i, pred in enumerate(predictions)])
+    return formatted_predictions  # Return formatted list of predictions
+
+# Custom Gradio interface with additional components
+def gradio_interface(file_path):
+    # Predict bird species
+    prediction = predict_bird(file_path)
+    
+    # Display the uploaded MP3 file with a play button
+    audio_player = gr.Audio(file_path, label="Uploaded MP3 File", visible=True, autoplay=True)
+    
+    # Display images with titles
+    bird_species_image = gr.Image("1.jpg", label="Bird Species")
+    bird_description_image = gr.Image("2.jpg", label="Bird Description")
+    bird_origins_image = gr.Image("3.jpg", label="Bird Origins")
+    
+    return prediction, audio_player, bird_species_image, bird_description_image, bird_origins_image
+
+# Launch Gradio interface
+interface = gr.Interface(
+    fn=gradio_interface,
+    inputs=[
+        gr.File(label="Upload MP3 file", file_types=['.mp3']),
+        gr.Audio(label="Uploaded MP3 File"),
+    ]
+    outputs=[
+        gr.Textbox(label="Predicted Bird Species"),
+        gr.Image(label="Bird Species"),
+        gr.Image(label="Bird Description"),
+        gr.Image(label="Bird Origins")
+    ]
+)