Update app.py

app.py

@@ -64,38 +64,45 @@ def get_last_conv_layer_name(model):
     return last_conv_layer
 
 # Function to generate Grad-CAM heatmap for a given beat and class index
-def make_gradcam_heatmap(beat, model, conv_layer_name, class_index):
-    # Create a model that maps the input beat to the activations of the conv layer and the output predictions
+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(
-        [model.inputs],
-        [model.get_layer(conv_layer_name).output, model.output]
+        inputs=model.inputs,
+        outputs=[model.get_layer(layer_name).output, model.output]
     )
-    # Record operations for automatic differentiation
+    
     with tf.GradientTape() as tape:
-        # Expand dims to add batch axis: shape (1, 257, 1)
-        beat_tensor = tf.expand_dims(beat, axis=0)
-        conv_outputs, predictions = grad_model(beat_tensor)
-        loss = predictions[:, class_index]
-    # Compute gradients of the target class wrt feature map
+        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 pooling over the time dimension to get weights
-    weights = tf.reduce_mean(grads, axis=1)
-    # Compute the weighted sum of feature maps along the channel dimension
-    cam = tf.reduce_sum(tf.multiply(weights, conv_outputs), axis=-1)
-    cam = tf.squeeze(cam)  # Remove batch dimension
-    # Apply ReLU to the heatmap to keep only positive influences
-    heatmap = tf.maximum(cam, 0)
-    # Normalize heatmap to the [0, 1] range
-    heatmap_max = tf.reduce_max(heatmap)
-    if heatmap_max == 0:
-        heatmap = tf.zeros_like(heatmap)
-    else:
-        heatmap /= heatmap_max
-    heatmap = heatmap.numpy()
-    # Resize heatmap to match the input beat size (if needed)
-    # For 1D, we use cv2.resize with the new shape (length, 1) then flatten
-    heatmap = cv2.resize(heatmap, (beat.shape[0], 1)).flatten()
-    return heatmap
+    
+    # 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")
@@ -227,36 +234,45 @@ if record_loaded and record is not None and annotation is not None:
     left_col, right_col = st.columns(2)
     
     def display_class_beats(col, class_name, beat_indices, num_beats):
-        """Helper function to display beats in a column"""
         with col:
             st.subheader(class_name)
             if len(beat_indices) == 0:
                 st.warning(f"No {class_name} beats found")
                 return
-            
-            num_to_show = min(num_beats, len(beat_indices))
-            for i, beat_idx in enumerate(beat_indices[:num_to_show]):
-                beat = beats[beat_idx]
-                pred_class = predicted_classes[beat_idx]
+
+            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 Grad-CAM heatmap
-                heatmap = make_gradcam_heatmap(beat, model, conv_layer_name, pred_class)
+                # generate the 1D heatmap
+                heatmap = generate_grad_cam(model, sample, conv_layer_name)
                 
-                # Create visualization
+                # set up figure
                 fig, ax = plt.subplots(figsize=(8, 2))
-                ax.plot(beat.flatten(), color="black")
-                sc = ax.scatter(
-                    np.arange(len(beat)),
-                    beat.flatten(),
-                    c=heatmap,
-                    cmap="jet",
-                    s=15
+                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}")
-                plt.colorbar(sc, ax=ax)
+                ax.set_xlim(0, len(beat))
+                ax.set_ylim(y_min, y_max)
+                
                 st.pyplot(fig)
-    
-    # Display left class beats
+        # Display left class beats
     display_class_beats(left_col, left_class, left_indices, num_beats)
     
     # Display right class beats