Cardio-AI_Assistant / visualization.py
MalikShehram's picture
Create visualization.py
5fc93e8 verified
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