app.py

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.")