Update app.py

app.py

@@ -113,8 +113,63 @@ def load_model_from_hub(species_type, model_version):
         print(f"Error loading {species_type} - {model_version}: {str(e)}")
         return None, None
 
-def predict_species(species_type, model_version, audio_file):
-    """Main prediction function"""
+def predict_with_model(spec, model, classes):
+    """Predict species from spectrogram - Adapted from notebook"""
+    # Ensure correct dimensions (1025, 313)
+    if spec.shape != (1025, 313):
+        # Resize if needed
+        spec = resize_spectrogram(spec, (1025, 313))
+    
+    # Preprocess for model - exactly as in notebook
+    arr = np.expand_dims(spec[..., np.newaxis], axis=0).astype('float32')
+    X = arr / np.max(arr)
+
+    # Predict
+    pred = model.predict(X, verbose=0)
+    pred_class_idx = np.argmax(pred)
+    pred_class = str(classes[pred_class_idx])
+    prob = float(pred[0][pred_class_idx])
+
+    return pred_class, prob
+
+def extract_chunks(audio_clean, sr, time=5):
+    """Extract audio chunks - Adapted from notebook's ext_chunks function"""
+    n_samples = sr * time
+    chunks = []
+    for i in range(0, len(audio_clean), n_samples):
+        start = i
+        end = i + n_samples
+
+        if end <= len(audio_clean):
+            chunk = audio_clean[start:end]
+        else:
+            # Circular padding - exactly as in notebook
+            faltan = end - len(audio_clean)
+            padding = audio_clean[:faltan]
+            chunk = np.concatenate([audio_clean[start:], padding])
+
+        chunks.append(chunk)
+    return np.array(chunks)
+
+def create_spectrogram(array_audio, n_fft=2048):
+    """Create spectrogram from audio array - Adapted from notebook's spectogram function"""
+    if isinstance(array_audio, np.ndarray):
+        dta = np.abs(librosa.stft(array_audio, n_fft=n_fft))
+        D = librosa.amplitude_to_db(dta, ref=np.max)
+    else:
+        dta = np.abs(librosa.stft(array_audio.numpy()))
+        D = librosa.amplitude_to_db(dta, ref=np.max)
+    return D
+
+def resize_spectrogram(spec, target_shape):
+    """Resize spectrogram to target shape"""
+    from scipy import ndimage
+    zoom_factors = (target_shape[0] / spec.shape[0], target_shape[1] / spec.shape[1])
+    resized = ndimage.zoom(spec, zoom_factors, order=1)
+    return resized
+
+def predict_species_all_chunks(species_type, model_version, audio_file):
+    """Main prediction function that processes all chunks"""
     global current_audio_path, current_audio_name
     
     if audio_file is None:
@@ -123,7 +178,6 @@ def predict_species(species_type, model_version, audio_file):
     try:
         # Store the current audio path and name for feedback
         current_audio_path = audio_file
-        # Extract the original audio file name
         current_audio_name = os.path.basename(audio_file)
         
         # Load model
@@ -131,24 +185,23 @@ def predict_species(species_type, model_version, audio_file):
         if model is None or classes is None:
             return pd.DataFrame({"Error": [f"Could not load {species_type} - {model_version} model"]})
         
-        # Process audio
+        # Process audio - using notebook approach
         wav, sr = torchaudio.load(audio_file)
         wav = wav.mean(dim=0)  # Convert to mono
         
-        # Extract 5-second chunks
+        # Extract 5-second chunks using notebook function
         chunks = extract_chunks(wav.numpy(), sr, time=5)
         
         results = []
         
         for i, chunk in enumerate(chunks):
-            # Create spectrogram
+            # Create spectrogram using notebook function
             spectrogram = create_spectrogram(chunk)
             
-            # Normalize
-            if np.std(spectrogram) > 0:
-                spectrogram = (spectrogram - np.mean(spectrogram)) / np.std(spectrogram)
+            # Normalize exactly as in notebook
+            spectrogram = (spectrogram - np.mean(spectrogram)) / np.std(spectrogram)
             
-            # Predict
+            # Predict using adapted notebook function
             species, confidence = predict_with_model(spectrogram, model, classes)
             
             time_start = i * 5
@@ -174,56 +227,85 @@ def predict_species(species_type, model_version, audio_file):
         print(f"Prediction error: {str(e)}")
         return pd.DataFrame({"Error": [f"Error during analysis: {str(e)}"]})
 
-def extract_chunks(audio_clean, sr, time):
-    """Extract audio chunks"""
-    n_samples = sr * time
-    chunks = []
-    for i in range(0, len(audio_clean), n_samples):
-        start = i
-        end = i + n_samples
-
-        if end <= len(audio_clean):
-            chunk = audio_clean[start:end]
-        else:
-            # Circular padding
-            missing = end - len(audio_clean)
-            padding = audio_clean[:missing]
-            chunk = np.concatenate([audio_clean[start:], padding])
-
-        chunks.append(chunk)
-    return np.array(chunks)
-
-def create_spectrogram(array_audio, n_fft=2048):
-    """Create spectrogram from audio array"""
-    dta = np.abs(librosa.stft(array_audio, n_fft=n_fft))
-    D = librosa.amplitude_to_db(dta, ref=np.max)
-    return D
-
-def predict_with_model(spec, model, classes):
-    """Predict species from spectrogram"""
-    # Ensure correct dimensions (1025, 313)
-    if spec.shape != (1025, 313):
-        # Resize if needed
-        spec = resize_spectrogram(spec, (1025, 313))
+def predict_species_final(species_type, model_version, audio_file):
+    """Enhanced prediction with voting system across chunks"""
+    global current_audio_path, current_audio_name
     
-    # Preprocess for model
-    arr = np.expand_dims(spec[..., np.newaxis], axis=0).astype('float32')
-    X = arr / np.max(arr)
-
-    # Predict
-    pred = model.predict(X, verbose=0)
-    pred_class_idx = np.argmax(pred)
-    pred_class = str(classes[pred_class_idx])
-    prob = float(pred[0][pred_class_idx])
-
-    return pred_class, prob
-
-def resize_spectrogram(spec, target_shape):
-    """Resize spectrogram to target shape"""
-    from scipy import ndimage
-    zoom_factors = (target_shape[0] / spec.shape[0], target_shape[1] / spec.shape[1])
-    resized = ndimage.zoom(spec, zoom_factors, order=1)
-    return resized
+    if audio_file is None:
+        return pd.DataFrame({"Info": ["Please upload an audio file"]})
+    
+    try:
+        current_audio_path = audio_file
+        current_audio_name = os.path.basename(audio_file)
+        
+        # Load model
+        model, classes = load_model_from_hub(species_type, model_version)
+        if model is None or classes is None:
+            return pd.DataFrame({"Error": [f"Could not load {species_type} - {model_version} model"]})
+        
+        # Process audio
+        wav, sr = torchaudio.load(audio_file)
+        wav = wav.mean(dim=0)
+        
+        # Extract chunks
+        chunks = extract_chunks(wav.numpy(), sr, time=5)
+        
+        results = []
+        species_votes = {}
+        
+        for i, chunk in enumerate(chunks):
+            # Create and normalize spectrogram
+            spectrogram = create_spectrogram(chunk)
+            spectrogram = (spectrogram - np.mean(spectrogram)) / np.std(spectrogram)
+            
+            # Predict
+            species, confidence = predict_with_model(spectrogram, model, classes)
+            
+            # Count votes for final prediction
+            if species in species_votes:
+                species_votes[species] += confidence
+            else:
+                species_votes[species] = confidence
+            
+            time_start = i * 5
+            time_end = (i + 1) * 5
+            
+            results.append({
+                'Segment': f'{i+1}',
+                'Time': f'{time_start}s - {time_end}s',
+                'Species': species,
+                'Confidence': f'{confidence:.1%}'
+            })
+        
+        # Determine final prediction
+        if species_votes:
+            final_species = max(species_votes, key=species_votes.get)
+            final_confidence = species_votes[final_species] / len(chunks)
+            
+            # Add final prediction row
+            final_row = pd.DataFrame({
+                'Segment': ['FINAL'],
+                'Time': ['Full Audio'],
+                'Species': [final_species],
+                'Confidence': [f'{final_confidence:.1%}']
+            })
+            
+            results_df = pd.concat([pd.DataFrame(results), final_row], ignore_index=True)
+        else:
+            results_df = pd.DataFrame(results)
+        
+        # Clean memory
+        del model
+        gc.collect()
+        
+        if results_df.empty:
+            return pd.DataFrame({"Info": ["No valid segments detected in the audio"]})
+        
+        return results_df
+        
+    except Exception as e:
+        print(f"Prediction error: {str(e)}")
+        return pd.DataFrame({"Error": [f"Error during analysis: {str(e)}"]})
 
 def save_feedback_to_dataset(audio_file_path, original_audio_name, feedback_text, consent_given, species_type, model_version, results_df):
     """Save audio and feedback to private Hugging Face dataset"""
@@ -341,6 +423,7 @@ with gr.Blocks(
     css="""
     .gradio-container { max-width: 1200px; margin: auto; }
     .consent-text { font-size: 0.9em; color: #666; }
+    .final-prediction { background-color: #e8f5e8 !important; font-weight: bold; }
     """
 ) as demo:
     
@@ -348,10 +431,16 @@ with gr.Blocks(
     current_results = gr.State(value=pd.DataFrame())
     
     gr.Markdown("""
-    #Species Audio Classifier
+    ## Species Audio Classifier
     **Upload an audio file to identify species using AI models**
     
-    *Models are loaded from: [RonaldCeballos/SpeciesClassifiers](https://huggingface.co/RonaldCeballos/SpeciesClassifiers)*
+    *Based on your notebook implementation - Models are loaded from: [RonaldCeballos/SpeciesClassifiers](https://huggingface.co/RonaldCeballos/SpeciesClassifiers)*
+    
+    🔍 **How it works:**
+    - Audio is split into 5-second segments
+    - Each segment is converted to a spectrogram
+    - AI model predicts species for each segment
+    - Final prediction is based on voting across all segments
     """)
     
     with gr.Row():
@@ -386,10 +475,11 @@ with gr.Blocks(
             gr.Markdown("""
             ### 💡 Instructions:
             1. Select species category
-            2. Choose model version
+            2. Choose model version  
             3. Upload audio file (WAV, MP3, etc.)
             4. Click "Analyze Audio"
             5. Review results by 5-second segments
+            6. Final prediction shown at the bottom
             """)
 
         with gr.Column(scale=2):
@@ -399,12 +489,14 @@ with gr.Blocks(
                 label="🎧 Analyzed Chunks",
                 headers=["Chunks", "Time", "Species", "Confidence"],
                 wrap=True,
-                max_height=400
+                max_height=500,
+                datatype=["str", "str", "str", "str"]
             )
             
             with gr.Accordion("💬 Submit Feedback for Model Improvement", open=False):
                 gr.Markdown("""
                 **Help us improve!** Submit your audio and feedback to our private dataset for model training.
+                *Using the same approach as your notebook implementation*
                 """)
                 
                 consent_checkbox = gr.Checkbox(
@@ -426,7 +518,7 @@ with gr.Blocks(
 
     # Event handlers
     predict_btn.click(
-        fn=predict_species,
+        fn=predict_species_final,  # Using the enhanced version with voting
         inputs=[species_selector, model_selector, audio_input],
         outputs=results_display
     ).then(