import streamlit as st import numpy as np import pandas as pd import tensorflow as tf import wfdb import tempfile import os from scipy.signal import resample import matplotlib.pyplot as plt import cv2 # Custom activation functions def sin_activation(x): return tf.math.sin(x) def cos_activation(x): return tf.math.cos(x) # Load model with custom objects @st.cache_resource def load_model(): return tf.keras.models.load_model( "model.keras", custom_objects={ 'sin': sin_activation, 'cos': cos_activation, 'gelu': tf.keras.activations.gelu } ) model = load_model() # AAMI class map class_map = { 0: "Normal", 1: "Supraventricular Ectopic (SVEB)", 2: "Ventricular Ectopic (VEB)", 3: "Fusion Beat", 4: "Unknown" } # Function to extract beats from record def extract_beats(record, annotation, window_size=257): beats = [] r_locs = annotation.sample signal = record.p_signal[:, 0] # Using first channel for r in r_locs: start = max(0, r - window_size//2) end = min(len(signal), r + window_size//2 + 1) if end - start == window_size: beat = signal[start:end] beats.append(beat) return np.array(beats) # Function to detect the last Conv1D layer in the model def get_last_conv_layer_name(model): last_conv_layer = None # Loop in reverse order over layers to find a Conv1D layer for layer in reversed(model.layers): if isinstance(layer, tf.keras.layers.Conv1D): last_conv_layer = layer.name break if last_conv_layer is None: st.error("No Conv1D layer found in the model. Grad-CAM requires a convolution layer.") return last_conv_layer # Function to generate Grad-CAM heatmap for a given beat and class index def generate_grad_cam(model, sample, layer_name): """ model : your loaded Keras model sample : a 4D tensor of shape (1, window_size, 1) layer_name : name of the Conv1D layer to use for Grad‑CAM returns : 1D numpy heatmap of length window_size """ # Build a model that returns both the conv outputs and the predictions grad_model = tf.keras.models.Model( inputs=model.inputs, outputs=[model.get_layer(layer_name).output, model.output] ) with tf.GradientTape() as tape: conv_outputs, predictions = grad_model(sample) # pick the top predicted class class_idx = tf.argmax(predictions[0]) loss = predictions[:, class_idx] # gradient of the loss wrt conv outputs grads = tape.gradient(loss, conv_outputs) # global average pool the gradients to get the importance of each channel pooled_grads = tf.reduce_mean(grads, axis=(0, 1)) # shape = (channels,) # remove batch dim from conv_outputs -> (time, channels) conv_outputs = tf.squeeze(conv_outputs, axis=0) # weight the conv outputs by the pooled gradients cam = tf.reduce_sum(conv_outputs * pooled_grads, axis=-1) # shape = (time,) raw = cam.numpy() print("raw min/max:", raw.min(), raw.max()) cam = tf.abs(cam) # ReLU cam = cam / (tf.reduce_max(cam) + 1e-8) # normalize return cam.numpy() # Streamlit App Layout st.title("ECG Arrhythmia Classification with Grad-CAM Visualization") st.write("Upload MIT-BIH record files (.dat, .hea, .atr) or load record 108") record_loaded = False record = None annotation = None # Load Record 108 Button if st.button("Load Record 108"): try: base_name = "108" record = wfdb.rdrecord(base_name) annotation = wfdb.rdann(base_name, 'atr') record_loaded = True except Exception as e: st.error(f"Error loading Record 108: {str(e)}") # File uploader uploaded_files = st.file_uploader( "Or upload your own files", type=["dat", "hea", "atr"], accept_multiple_files=True ) if uploaded_files and not record_loaded: with tempfile.TemporaryDirectory() as tmpdir: for f in uploaded_files: file_path = os.path.join(tmpdir, f.name) with open(file_path, "wb") as f_out: f_out.write(f.getbuffer()) base_names = {os.path.splitext(f.name)[0] for f in uploaded_files} common_base = list(base_names)[0] try: record = wfdb.rdrecord(os.path.join(tmpdir, common_base)) annotation = wfdb.rdann(os.path.join(tmpdir, common_base), 'atr') record_loaded = True except Exception as e: st.error(f"Error reading uploaded files: {str(e)}") # Process the record if loaded if record_loaded and record is not None and annotation is not None: beats = extract_beats(record, annotation) if len(beats) == 0: st.error("No valid beats found in the record") st.stop() beats = beats.reshape((-1, 257, 1)).astype(np.float32) predictions = model.predict(beats) predicted_classes = np.argmax(predictions, axis=1) st.subheader("Classification Results") results = pd.DataFrame({ "Beat Index": range(len(beats)), "Predicted Class": [class_map[c] for c in predicted_classes], "Confidence": np.max(predictions, axis=1) }) st.dataframe(results) # Class Distribution Section st.subheader("Class Distribution") class_indices = list(class_map.keys()) class_names = [class_map[i] for i in class_indices] counts = [np.sum(predicted_classes == i) for i in class_indices] distribution_df = pd.DataFrame({ "Class": class_names, "Count": counts }) col1, col2 = st.columns([1, 2]) with col1: st.dataframe(distribution_df.style.format({'Count': '{:,}'})) with col2: st.bar_chart(distribution_df.set_index('Class')) # Display a Sample ECG Beat st.subheader("Sample ECG Beat") fig, ax = plt.subplots() ax.plot(beats[0].flatten(), label="ECG Beat") ax.legend() st.pyplot(fig) # ---------------- Grad-CAM Visualization Section ---------------- st.subheader("Class Comparison with Grad-CAM") st.write("Compare model explanations between classes present in this record") # Automatically detect the last convolutional layer name conv_layer_name = get_last_conv_layer_name(model) if conv_layer_name is not None: st.write(f"Using Conv1D layer: **{conv_layer_name}** for Grad-CAM.") # Get classes actually present in the data present_classes = distribution_df[distribution_df['Count'] > 0]['Class'].tolist() if not present_classes: st.warning("No classes with detected beats to compare") st.stop() # Class selection dropdowns col1, col2, col3 = st.columns([1, 1, 1]) with col1: left_class = st.selectbox( "Left Class:", options=present_classes, index=0 ) with col2: # Default to second class if available, else first right_index = 1 if len(present_classes) > 1 else 0 right_class = st.selectbox( "Right Class:", options=present_classes, index=right_index ) with col3: num_beats = st.number_input( "Beats per class:", min_value=1, max_value=10, value=3 ) # Get class indices from names class_name_to_idx = {v: k for k, v in class_map.items()} left_class_idx = class_name_to_idx[left_class] right_class_idx = class_name_to_idx[right_class] left_indices = np.where(predicted_classes == left_class_idx)[0] right_indices = np.where(predicted_classes == right_class_idx)[0] # Create comparison columns left_col, right_col = st.columns(2) def display_class_beats(col, class_name, beat_indices, num_beats): with col: st.subheader(class_name) if len(beat_indices) == 0: st.warning(f"No {class_name} beats found") return for beat_idx in beat_indices[:num_beats]: beat = beats[beat_idx].flatten() # shape (window_size,) sample = beat.reshape(1, -1, 1).astype(np.float32) # generate the 1D heatmap heatmap = generate_grad_cam(model, sample, conv_layer_name) # set up figure fig, ax = plt.subplots(figsize=(8, 2)) y_min, y_max = beat.min(), beat.max() # Always draw the heatmap background for all beats ax.imshow( np.expand_dims(heatmap, axis=0), # shape (1, window_size) aspect='auto', cmap='jet', alpha=0.5, extent=[0, len(beat), y_min, y_max] ) # overlay the ECG trace ax.plot(beat, linewidth=2, color='blue') # styling # Do NOT set a facecolor here - it will block the heatmap # ax.set_facecolor('#e0e0f0') # This line is commented out ax.axis('off') # clean look ax.set_title(f"Beat {beat_idx}") ax.set_xlim(0, len(beat)) ax.set_ylim(y_min, y_max) st.pyplot(fig) # Display left class beats display_class_beats(left_col, left_class, left_indices, num_beats) # Display right class beats display_class_beats(right_col, right_class, right_indices, num_beats) # Add comparison note if same class selected if left_class == right_class: st.info("Comparing different instances of the same class. Note: This shows intra-class variation.")