Create visualization.py

visualization.py

@@ -0,0 +1,68 @@
+import librosa
+import librosa.display
+import matplotlib.pyplot as plt
+import numpy as np
+
+def generate_visualizations(audio_path):
+    """
+    Takes an audio file path, processes it, and returns two matplotlib figures:
+    1. A Time-Domain PCG Waveform
+    2. A Frequency-Domain Mel-Spectrogram tuned for heart sounds.
+    """
+    if not audio_path:
+        return None, None
+
+    try:
+        # Load the audio file
+        # sr=None preserves the original sample rate. 
+        y, sr = librosa.load(audio_path, sr=None)
+        
+        # --- 1. Generate PCG Waveform (Time Domain) ---
+        # Close any existing plots to prevent memory leaks in Gradio
+        plt.close('all') 
+        
+        fig_pcg, ax_pcg = plt.subplots(figsize=(10, 4))
+        librosa.display.waveshow(y, sr=sr, ax=ax_pcg, color="#1f77b4", alpha=0.8)
+        
+        ax_pcg.set_title("Phonocardiogram (PCG) Waveform", fontsize=14, fontweight="bold")
+        ax_pcg.set_xlabel("Time (seconds)")
+        ax_pcg.set_ylabel("Amplitude")
+        ax_pcg.grid(True, linestyle="--", alpha=0.6)
+        fig_pcg.tight_layout()
+
+        # --- 2. Generate Mel-Spectrogram (Frequency Domain) ---
+        fig_spec, ax_spec = plt.subplots(figsize=(10, 4))
+        
+        # Compute the Mel-spectrogram
+        # We set fmax=2000 because heart sounds and murmurs rarely exceed 2000 Hz.
+        # This focuses the graph on the relevant medical data.
+        S = librosa.feature.melspectrogram(y=y, sr=sr, n_mels=128, fmax=2000)
+        
+        # Convert power to decibels (log scale) for better visualization
+        S_dB = librosa.power_to_db(S, ref=np.max)
+        
+        # Display the spectrogram
+        img = librosa.display.specshow(
+            S_dB, 
+            x_axis='time', 
+            y_axis='mel', 
+            sr=sr, 
+            fmax=2000, 
+            ax=ax_spec, 
+            cmap='magma' # 'magma' or 'viridis' looks very professional for medical imaging
+        )
+        
+        ax_spec.set_title("Mel-Spectrogram (Low-Frequency Focus)", fontsize=14, fontweight="bold")
+        ax_spec.set_xlabel("Time (seconds)")
+        ax_spec.set_ylabel("Frequency (Hz)")
+        fig_spec.colorbar(img, ax=ax_spec, format='%+2.0f dB', label="Intensity (dB)")
+        fig_spec.tight_layout()
+
+        return fig_pcg, fig_spec
+
+    except Exception as e:
+        print(f"Error generating visualizations: {e}")
+        # Return empty figures in case of an error so the UI doesn't crash
+        fig, ax = plt.subplots()
+        ax.text(0.5, 0.5, "Error generating graph", ha="center")
+        return fig, fig