Update app.py

app.py

@@ -178,39 +178,90 @@ if record_loaded and record is not None and annotation is not None:
     st.pyplot(fig)
 
     # ---------------- Grad-CAM Visualization Section ----------------
-    st.subheader("Grad-CAM Heatmap Visualization for Each Beat")
-    st.write("Below are Grad-CAM heatmaps overlaying each beat. The heatmaps show the regions contributing most to the predicted class.")
+    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.")
-        
-        # Optionally, you can limit the number of beats displayed to avoid long processing times.
-        # For demonstration, here we process all beats, but you might want to show only the first N beats.
-        show_all = st.checkbox("Show Grad-CAM for all beats", value=False)
-        if not show_all:
-            num_beats_to_show = st.number_input("Number of beats to show:", min_value=1, max_value=len(beats), value=5)
-        else:
-            num_beats_to_show = len(beats)
-        
-        # Loop over each beat and its prediction to generate Grad-CAM heatmap
-        for idx in range(num_beats_to_show):
-            beat = beats[idx]
-            pred_class = predicted_classes[idx]
-            predicted_label = class_map[pred_class]
-            # Compute Grad-CAM heatmap for the beat
-            heatmap = make_gradcam_heatmap(beat, model, conv_layer_name, pred_class)
+    
+    # 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):
+        """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
             
-            # Generate visualization figure
-            fig, ax = plt.subplots(figsize=(10, 3))
-            # Plot the raw ECG beat
-            ax.plot(beat.flatten(), color="black", label="ECG Beat")
-            # Overlay Grad-CAM heatmap by scatter plotting points with a colormap according to heatmap value
-            sc = ax.scatter(np.arange(len(beat)), beat.flatten(), c=heatmap, cmap="jet", s=25)
-            ax.set_title(f"Beat {idx} - Predicted: {predicted_label}")
-            ax.set_xlabel("Time Index")
-            ax.set_ylabel("Amplitude")
-            # Add a colorbar to indicate heatmap intensity
-            fig.colorbar(sc, ax=ax, label="Grad-CAM Intensity")
-            st.pyplot(fig)
+            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]
+                
+                # Generate Grad-CAM heatmap
+                heatmap = make_gradcam_heatmap(beat, model, conv_layer_name, pred_class)
+                
+                # Create visualization
+                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
+                )
+                ax.set_title(f"Beat {beat_idx}")
+                plt.colorbar(sc, ax=ax)
+                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.")