Update app.py

app.py

@@ -54,7 +54,6 @@ def extract_beats(record, annotation, window_size=257):
 # 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
@@ -65,13 +64,6 @@ def get_last_conv_layer_name(model):
 
 # 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]
@@ -79,45 +71,40 @@ def generate_grad_cam(model, sample, layer_name):
     
     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)
+    pooled_grads = tf.reduce_mean(grads, axis=(0, 1))
     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,)
+    cam = tf.reduce_sum(conv_outputs * pooled_grads, axis=-1)
     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
+    cam = tf.abs(cam)
+    cam = cam / (tf.reduce_max(cam) + 1e-8)
     
     return cam.numpy()
 
-
+# Initialize session state variables if not already set
+if 'record_loaded' not in st.session_state:
+    st.session_state.record_loaded = False
+if 'record' not in st.session_state:
+    st.session_state.record = None
+if 'annotation' not in st.session_state:
+    st.session_state.annotation = None
 
 # 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
+        st.session_state.record = wfdb.rdrecord(base_name)
+        st.session_state.annotation = wfdb.rdann(base_name, 'atr')
+        st.session_state.record_loaded = True
     except Exception as e:
         st.error(f"Error loading Record 108: {str(e)}")
 
@@ -128,7 +115,7 @@ uploaded_files = st.file_uploader(
     accept_multiple_files=True
 )
 
-if uploaded_files and not record_loaded:
+if uploaded_files and not st.session_state.record_loaded:
     with tempfile.TemporaryDirectory() as tmpdir:
         for f in uploaded_files:
             file_path = os.path.join(tmpdir, f.name)
@@ -137,15 +124,15 @@ if uploaded_files and not record_loaded:
         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
+            st.session_state.record = wfdb.rdrecord(os.path.join(tmpdir, common_base))
+            st.session_state.annotation = wfdb.rdann(os.path.join(tmpdir, common_base), 'atr')
+            st.session_state.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 st.session_state.record_loaded and st.session_state.record is not None and st.session_state.annotation is not None:
+    beats = extract_beats(st.session_state.record, st.session_state.annotation)
     if len(beats) == 0:
         st.error("No valid beats found in the record")
         st.stop()
@@ -162,7 +149,6 @@ if record_loaded and record is not None and annotation is not None:
     })
     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]
@@ -177,62 +163,40 @@ if record_loaded and record is not None and annotation is not None:
     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
-        )
+        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
-        )
+        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
-        )
+        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)
@@ -241,43 +205,26 @@ if record_loaded and record is not None and annotation is not None:
                 return
 
             for beat_idx in beat_indices[:num_beats]:
-                beat = beats[beat_idx].flatten()             # shape (window_size,)
+                beat = beats[beat_idx].flatten()
                 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)
+                    np.expand_dims(heatmap, axis=0),
                     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.axis('off')
                 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.")