Upload 5 files

app.py

@@ -190,60 +190,102 @@ def predict(image):
 
 def make_gradcam_heatmap(img_array, model, pred_index=None):
     """
-    Generate GradCAM heatmap for the model prediction.
+    Generate Grad-CAM heatmap for lightweight ViT model
+    Uses the transformer output before global pooling
     """
     try:
-        # Find the last convolutional layer
-        last_conv_layer = None
+        # Find the layer before GlobalAveragePooling (typically the last Add or LayerNormalization)
+        target_layer = None
         for layer in reversed(model.layers):
-            if 'conv' in layer.name.lower() or isinstance(layer, tf.keras.layers.Conv2D):
-                last_conv_layer = layer
+            # Look for the last Add layer (from transformer blocks)
+            if isinstance(layer, tf.keras.layers.Add):
+                target_layer = layer
+                break
+            # Or the LayerNormalization before classification head
+            if isinstance(layer, tf.keras.layers.LayerNormalization) and 'representation' not in layer.name:
+                target_layer = layer
                 break
 
-        if last_conv_layer is None:
-            # If no conv layer, try to use patch encoder or dense layers
+        if target_layer is None:
+            # Fallback: find any layer with 3D output (batch, seq_len, features)
             for layer in reversed(model.layers):
-                if 'dense' in layer.name.lower() or 'patch' in layer.name.lower():
-                    last_conv_layer = layer
+                if hasattr(layer, 'output_shape') and len(layer.output_shape) == 3:
+                    target_layer = layer
                     break
 
-        if last_conv_layer is None:
-            return None
+        if target_layer is None:
+            print("Warning: No suitable layer found for Grad-CAM")
+            return None, pred_index
 
-        # Create a model that maps the input to the activations of the last conv layer and the output predictions
+        # Create a model that outputs both the target layer output and final predictions
         grad_model = tf.keras.models.Model(
-            [model.inputs],
-            [last_conv_layer.output, model.output]
+            inputs=model.inputs,
+            outputs=[model.get_layer(target_layer.name).output, model.output]
         )
 
-        # Compute the gradient of the top predicted class with respect to the output feature map
+        # Compute gradients
         with tf.GradientTape() as tape:
-            conv_outputs, predictions = grad_model(img_array)
+            layer_output, predictions = grad_model(img_array, training=False)
             if pred_index is None:
                 pred_index = tf.argmax(predictions[0])
             class_channel = predictions[:, pred_index]
 
-        # Gradient of the output neuron with respect to the output feature map
-        grads = tape.gradient(class_channel, conv_outputs)
-
-        # Vector of mean intensity of the gradient over a specific feature map channel
-        pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
-
-        # Multiply each channel in the feature map array by importance
-        conv_outputs = conv_outputs[0]
-        heatmap = conv_outputs @ pooled_grads[..., tf.newaxis]
-        heatmap = tf.squeeze(heatmap)
+        # Get gradients of the predicted class with respect to the layer output
+        grads = tape.gradient(class_channel, layer_output)
+
+        if grads is None:
+            print("Warning: Gradients are None. Using simple attention map.")
+            # Fallback: use attention weights
+            layer_output_np = layer_output[0].numpy()
+            heatmap = np.mean(np.abs(layer_output_np), axis=-1)
+            # Reshape to 2D grid
+            num_patches = heatmap.shape[0]
+            grid_size = int(np.sqrt(num_patches))
+            heatmap = heatmap[:grid_size *
+                              grid_size].reshape(grid_size, grid_size)
+            heatmap = (heatmap - heatmap.min()) / \
+                (heatmap.max() - heatmap.min() + 1e-10)
+            return heatmap, int(pred_index.numpy())
+
+        # Global average pooling on gradients
+        if len(grads.shape) == 3:  # (batch, seq_len, features)
+            pooled_grads = tf.reduce_mean(grads, axis=(0, 1))
+            layer_output = layer_output[0]
+
+            # Weight the sequence by the gradients
+            heatmap = layer_output @ pooled_grads[..., tf.newaxis]
+            heatmap = tf.squeeze(heatmap)
+
+            # Reshape to 2D grid
+            num_patches = heatmap.shape[0]
+            grid_size = int(np.sqrt(num_patches))
+            if grid_size * grid_size != num_patches:
+                # Handle case where sqrt is not exact
+                # Exclude class token if present
+                grid_size = int(np.sqrt(num_patches - 1))
+                heatmap = heatmap[1:grid_size*grid_size+1]  # Skip class token
+            else:
+                heatmap = heatmap[:grid_size*grid_size]
+            heatmap = tf.reshape(heatmap, (grid_size, grid_size))
+        else:
+            pooled_grads = tf.reduce_mean(grads, axis=(0, 1, 2))
+            layer_output = layer_output[0]
+            heatmap = layer_output @ pooled_grads[..., tf.newaxis]
+            heatmap = tf.squeeze(heatmap)
 
-        # Normalize the heatmap
-        heatmap = tf.maximum(heatmap, 0) / tf.math.reduce_max(heatmap)
-        return heatmap.numpy()
+        # Normalize between 0 and 1
+        heatmap = tf.maximum(heatmap, 0) / \
+            (tf.math.reduce_max(heatmap) + 1e-10)
+        return heatmap.numpy(), int(pred_index.numpy())
 
     except Exception as e:
         print(f"GradCAM error: {e}")
-        return None
+        import traceback
+        traceback.print_exc()
+        return None, pred_index
 
 
-def apply_gradcam(image, heatmap):
+def apply_gradcam(image, heatmap, alpha=0.4):
     """
     Apply GradCAM heatmap overlay on the original image.
     """
@@ -251,21 +293,30 @@ def apply_gradcam(image, heatmap):
         if heatmap is None:
             return image
 
-        # Resize heatmap to match input image size
-        heatmap = cv2.resize(heatmap, (image_size, image_size))
-
-        # Convert heatmap to RGB
-        heatmap = np.uint8(255 * heatmap)
-        heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
-
         # Convert image to numpy array
         if isinstance(image, Image.Image):
             img_array = np.array(image.resize((image_size, image_size)))
         else:
             img_array = image
 
+        # Resize heatmap to match input image size
+        heatmap_resized = cv2.resize(
+            heatmap, (img_array.shape[1], img_array.shape[0]))
+
+        # Convert heatmap to RGB
+        heatmap_uint8 = np.uint8(255 * heatmap_resized)
+        heatmap_colored = cv2.applyColorMap(heatmap_uint8, cv2.COLORMAP_JET)
+        heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)
+
+        # Normalize image if needed
+        if img_array.max() <= 1.0:
+            img_uint8 = (img_array * 255).astype('uint8')
+        else:
+            img_uint8 = img_array.astype('uint8')
+
         # Superimpose the heatmap on original image
-        superimposed_img = cv2.addWeighted(img_array, 0.6, heatmap, 0.4, 0)
+        superimposed_img = heatmap_colored * alpha + img_uint8 * (1 - alpha)
+        superimposed_img = np.clip(superimposed_img, 0, 255).astype('uint8')
 
         return Image.fromarray(superimposed_img)
 
@@ -279,7 +330,7 @@ def generate_gradcam(image):
     Generate GradCAM visualization.
     """
     if model is None or image is None:
-        return None, "Model not loaded or no image provided"
+        return None
 
     try:
         # Preprocess image
@@ -288,23 +339,21 @@ def generate_gradcam(image):
         # Make prediction
         predictions = model.predict(processed_image, verbose=0)
         pred_class = np.argmax(predictions[0])
-        confidence = predictions[0][pred_class]
 
         # Generate heatmap
-        heatmap = make_gradcam_heatmap(processed_image, model, pred_class)
+        heatmap, _ = make_gradcam_heatmap(processed_image, model, pred_class)
 
         if heatmap is None:
-            return None, "GradCAM not available for this model architecture"
+            return None
 
         # Apply heatmap
-        gradcam_image = apply_gradcam(image, heatmap)
-
-        info = f"**Predicted Class:** {class_names[pred_class]}\n**Confidence:** {confidence:.2%}"
+        gradcam_image = apply_gradcam(image, heatmap, alpha=0.4)
 
-        return gradcam_image, info
+        return gradcam_image
 
     except Exception as e:
-        return None, f"Error generating GradCAM: {str(e)}"
+        print(f"Error generating GradCAM: {e}")
+        return None
 
 
 # -----------------------------
@@ -317,10 +366,10 @@ def generate_shap(image):
     Generate SHAP explanation visualization.
     """
     if not SHAP_AVAILABLE:
-        return None, "SHAP library not available. Please install shap package."
+        return None
 
     if model is None or image is None:
-        return None, "Model not loaded or no image provided"
+        return None
 
     try:
         # Preprocess image
@@ -362,18 +411,11 @@ def generate_shap(image):
         shap_image = Image.open(buf)
         plt.close()
 
-        # Get prediction info
-        prediction = model_predict(img_array[np.newaxis, ...])[0]
-        pred_class = np.argmax(prediction)
-        confidence = prediction[pred_class]
-
-        info = f"**Predicted Class:** {class_names[pred_class]}\n**Confidence:** {confidence:.2%}"
-
-        return shap_image, info
+        return shap_image
 
     except Exception as e:
         print(f"SHAP error: {e}")
-        return None, f"Error generating SHAP: {str(e)}"
+        return None
 
 
 # -----------------------------
@@ -386,10 +428,10 @@ def generate_lime(image):
     Generate LIME explanation visualization.
     """
     if not LIME_AVAILABLE or not SKIMAGE_AVAILABLE:
-        return None, None, "LIME or scikit-image library not available. Please install lime and scikit-image packages."
+        return None, None
 
     if model is None or image is None:
-        return None, None, "Model not loaded or no image provided"
+        return None, None
 
     try:
         # Preprocess image
@@ -414,12 +456,6 @@ def generate_lime(image):
             batch_size=32
         )
 
-        # Get prediction
-        prediction = model.predict(
-            img_normalized.reshape(1, image_size, image_size, 3))
-        pred_class = np.argmax(prediction)
-        confidence = np.max(prediction)
-
         # Create visualizations
         # Positive features only
         temp_positive, mask_positive = explanation.get_image_and_mask(
@@ -444,13 +480,43 @@ def generate_lime(image):
             (lime_positive * 255).astype(np.uint8))
         lime_both_img = Image.fromarray((lime_both * 255).astype(np.uint8))
 
-        info = f"**Predicted Class:** {class_names[pred_class]}\n**Confidence:** {confidence:.2%}"
-
-        return lime_positive_img, lime_both_img, info
+        return lime_positive_img, lime_both_img
 
     except Exception as e:
         print(f"LIME error: {e}")
-        return None, None, f"Error generating LIME: {str(e)}"
+        return None, None
+
+
+# -----------------------------
+# Unified Prediction with XAI
+# -----------------------------
+
+
+def predict_with_xai(image):
+    """
+    Make prediction and generate all XAI explanations at once.
+    """
+    if model is None or image is None:
+        return {class_name: 0.0 for class_name in class_names}, None, None, None, None
+
+    try:
+        # Make prediction
+        prediction_results = predict(image)
+
+        # Generate GradCAM
+        gradcam_img = generate_gradcam(image)
+
+        # Generate SHAP (can be slow)
+        shap_img = generate_shap(image)
+
+        # Generate LIME (can be slow)
+        lime_positive, lime_both = generate_lime(image)
+
+        return prediction_results, gradcam_img, shap_img, lime_positive, lime_both
+
+    except Exception as e:
+        print(f"Error in predict_with_xai: {e}")
+        return {class_name: 0.0 for class_name in class_names}, None, None, None, None
 
 
 # -----------------------------
@@ -461,8 +527,8 @@ description = """
 <div style="text-align: center; padding: 20px;">
     <h2 style="color: #2E86AB;">Advanced Medical Image Analysis with Explainable AI</h2>
     <p style="font-size: 16px; color: #555;">
-        Upload an endoscopic image to classify using a <b>Lightweight Vision Transformer</b> model
-        and explore <b>multiple explainability methods</b> to understand the model's decision-making process.
+        Upload an endoscopic image to classify using a <b>Lightweight Vision Transformer</b> model.
+        Get predictions with <b>three explainability methods</b> to understand the AI's decision.
     </p>
     <div style="display: flex; justify-content: center; gap: 20px; margin-top: 15px;">
         <div style="background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); padding: 15px; border-radius: 10px; color: white;">
@@ -484,20 +550,6 @@ custom_css = """
 .gradio-container {
     font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif !important;
 }
-.tab-nav button {
-    font-size: 16px !important;
-    font-weight: bold !important;
-    padding: 12px 24px !important;
-    border-radius: 10px 10px 0 0 !important;
-    transition: all 0.3s ease !important;
-}
-.tab-nav button.selected {
-    background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
-    color: white !important;
-}
-.tab-nav button:hover {
-    transform: translateY(-2px) !important;
-}
 h1 {
     background: linear-gradient(135deg, #667eea 0%, #764ba2 100%);
     -webkit-background-clip: text;
@@ -505,12 +557,6 @@ h1 {
     font-size: 2.5em !important;
     text-align: center !important;
 }
-.output-class {
-    font-size: 18px !important;
-    padding: 10px !important;
-    border-radius: 8px !important;
-    background: linear-gradient(135deg, #e0f7fa 0%, #e1bee7 100%) !important;
-}
 .button-primary {
     background: linear-gradient(135deg, #667eea 0%, #764ba2 100%) !important;
     border: none !important;
@@ -518,6 +564,7 @@ h1 {
     font-weight: bold !important;
     padding: 12px 30px !important;
     border-radius: 25px !important;
+    font-size: 16px !important;
     transition: all 0.3s ease !important;
 }
 .button-primary:hover {
@@ -526,103 +573,98 @@ h1 {
 }
 """
 
-examples = []  # Add example image paths if available
-
 # Create Gradio interface using Blocks with creative design
 with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
     gr.HTML(f"<h1>{title}</h1>")
     gr.HTML(description)
 
-    with gr.Tabs():
-        # Tab 1: Basic Prediction
-        with gr.Tab("📊 Classification"):
-            with gr.Row():
-                with gr.Column(scale=1):
-                    input_image_basic = gr.Image(
-                        type="pil", label="Upload Endoscopic Image")
-                    classify_btn = gr.Button(
-                        "🔍 Classify Image", variant="primary", elem_classes="button-primary")
-                    gr.Markdown("""
-                    <div style="background: #f0f4f8; padding: 15px; border-radius: 10px; margin-top: 10px;">
-                        <b>ℹ️ Instructions:</b>
-                        <ul>
-                            <li>Upload an endoscopic image (JPG, PNG)</li>
-                            <li>Click "Classify Image" to get predictions</li>
-                            <li>View confidence scores for each class</li>
-                        </ul>
-                    </div>
-                    """)
-
-                with gr.Column(scale=1):
-                    output_label_basic = gr.Label(
-                        num_top_classes=4, label="📈 Prediction Results")
-
-        # Tab 2: GradCAM Explanation
-        with gr.Tab("🔥 GradCAM"):
-            with gr.Row():
-                with gr.Column(scale=1):
-                    input_image_gradcam = gr.Image(
-                        type="pil", label="Upload Endoscopic Image")
-                    gradcam_btn = gr.Button(
-                        "🔥 Generate GradCAM", variant="primary", elem_classes="button-primary")
-                    gr.Markdown("""
-                    <div style="background: #fff3e0; padding: 15px; border-radius: 10px; margin-top: 10px;">
-                        <b>🔥 GradCAM (Gradient-weighted Class Activation Mapping):</b>
-                        <p>Highlights the regions of the image that are most important for the model's prediction.
-                        Red areas indicate high importance, while blue areas indicate low importance.</p>
-                    </div>
-                    """)
-
-                with gr.Column(scale=1):
-                    output_gradcam = gr.Image(label="🗺️ GradCAM Heatmap")
-                    output_gradcam_info = gr.Markdown(label="Prediction Info")
-
-        # Tab 3: SHAP Explanation
-        with gr.Tab("🎯 SHAP"):
-            with gr.Row():
-                with gr.Column(scale=1):
-                    input_image_shap = gr.Image(
-                        type="pil", label="Upload Endoscopic Image")
-                    shap_btn = gr.Button(
-                        "🎯 Generate SHAP", variant="primary", elem_classes="button-primary")
-                    gr.Markdown("""
-                    <div style="background: #e8f5e9; padding: 15px; border-radius: 10px; margin-top: 10px;">
-                        <b>🎯 SHAP (SHapley Additive exPlanations):</b>
-                        <p>Shows how each pixel contributes to the prediction. Red pixels push the prediction 
-                        towards the predicted class, while blue pixels push it away.</p>
-                        <p><i>⚠️ Note: SHAP generation may take 30-60 seconds.</i></p>
-                    </div>
-                    """)
-
-                with gr.Column(scale=1):
-                    output_shap = gr.Image(label="📊 SHAP Explanation")
-                    output_shap_info = gr.Markdown(label="Prediction Info")
-
-        # Tab 4: LIME Explanation
-        with gr.Tab("🍋 LIME"):
-            with gr.Row():
-                with gr.Column(scale=1):
-                    input_image_lime = gr.Image(
-                        type="pil", label="Upload Endoscopic Image")
-                    lime_btn = gr.Button(
-                        "🍋 Generate LIME", variant="primary", elem_classes="button-primary")
-                    gr.Markdown("""
-                    <div style="background: #fce4ec; padding: 15px; border-radius: 10px; margin-top: 10px;">
-                        <b>🍋 LIME (Local Interpretable Model-agnostic Explanations):</b>
-                        <p>Segments the image into superpixels and shows which segments are most important.
-                        Green boundaries indicate positive contributions to the prediction.</p>
-                        <p><i>⚠️ Note: LIME generation may take 30-60 seconds.</i></p>
-                    </div>
-                    """)
-
-                with gr.Column(scale=1):
-                    gr.Markdown("### 🟢 Positive Features Only")
-                    output_lime_positive = gr.Image(
-                        label="Important Regions (Positive)")
-                    gr.Markdown("### 🔴 Positive & Negative Features")
-                    output_lime_both = gr.Image(
-                        label="All Contributing Regions")
-                    output_lime_info = gr.Markdown(label="Prediction Info")
+    with gr.Row():
+        with gr.Column(scale=1):
+            input_image = gr.Image(
+                type="pil", label="📤 Upload Endoscopic Image")
+            predict_btn = gr.Button(
+                "🔍 Classify & Explain", variant="primary", elem_classes="button-primary", size="lg")
+
+            gr.Markdown("""
+            <div style="background: #f0f4f8; padding: 15px; border-radius: 10px; margin-top: 10px;">
+                <b>ℹ️ Instructions:</b>
+                <ul>
+                    <li>Upload an endoscopic image (JPG, PNG)</li>
+                    <li>Click "Classify & Explain" to get results</li>
+                    <li>View prediction + XAI explanations below</li>
+                    <li><i>Note: SHAP and LIME may take 30-60 seconds</i></li>
+                </ul>
+            </div>
+            """)
+
+        with gr.Column(scale=1):
+            output_label = gr.Label(
+                num_top_classes=4, label="📊 Prediction Results", show_label=True)
+
+    # Explanations Section
+    gr.Markdown("""
+    <div style="text-align: center; margin-top: 30px; margin-bottom: 20px;">
+        <h2 style="color: #2E86AB;">🎯 Explainable AI Visualizations</h2>
+        <p style="color: #666;">Understanding how the model makes its predictions</p>
+    </div>
+    """)
+
+    with gr.Row():
+        # GradCAM
+        with gr.Column(scale=1):
+            gr.Markdown("""
+            <div style="background: linear-gradient(135deg, #fff3e0 0%, #ffe0b2 100%); padding: 15px; border-radius: 10px; margin-bottom: 10px;">
+                <h3 style="margin: 0; color: #e65100;">🔥 Grad-CAM</h3>
+                <p style="margin: 5px 0 0 0; font-size: 14px;">
+                    <b>Gradient-weighted Class Activation Mapping</b><br>
+                    Highlights regions most important for prediction. Red = high importance.
+                </p>
+            </div>
+            """)
+            output_gradcam = gr.Image(
+                label="Grad-CAM Heatmap", show_label=False)
+
+    with gr.Row():
+        # SHAP
+        with gr.Column(scale=1):
+            gr.Markdown("""
+            <div style="background: linear-gradient(135deg, #e8f5e9 0%, #c8e6c9 100%); padding: 15px; border-radius: 10px; margin-bottom: 10px;">
+                <h3 style="margin: 0; color: #2e7d32;">🎯 SHAP</h3>
+                <p style="margin: 5px 0 0 0; font-size: 14px;">
+                    <b>SHapley Additive exPlanations</b><br>
+                    Red pixels push toward predicted class, blue pixels push away.
+                </p>
+            </div>
+            """)
+            output_shap = gr.Image(label="SHAP Explanation", show_label=False)
+
+    with gr.Row():
+        # LIME
+        with gr.Column(scale=1):
+            gr.Markdown("""
+            <div style="background: linear-gradient(135deg, #fce4ec 0%, #f8bbd0 100%); padding: 15px; border-radius: 10px; margin-bottom: 10px;">
+                <h3 style="margin: 0; color: #c2185b;">🍋 LIME - Positive Features</h3>
+                <p style="margin: 5px 0 0 0; font-size: 14px;">
+                    <b>Local Interpretable Model-agnostic Explanations</b><br>
+                    Green boundaries show regions supporting the prediction.
+                </p>
+            </div>
+            """)
+            output_lime_positive = gr.Image(
+                label="LIME Positive", show_label=False)
+
+        with gr.Column(scale=1):
+            gr.Markdown("""
+            <div style="background: linear-gradient(135deg, #e1f5fe 0%, #b3e5fc 100%); padding: 15px; border-radius: 10px; margin-bottom: 10px;">
+                <h3 style="margin: 0; color: #01579b;">🍋 LIME - All Features</h3>
+                <p style="margin: 5px 0 0 0; font-size: 14px;">
+                    <b>Positive & Negative Contributions</b><br>
+                    Shows both supporting and opposing regions.
+                </p>
+            </div>
+            """)
+            output_lime_both = gr.Image(
+                label="LIME Positive & Negative", show_label=False)
 
     # Footer
     gr.Markdown("""
@@ -636,15 +678,14 @@ with gr.Blocks(css=custom_css, theme=gr.themes.Soft()) as demo:
     </div>
     """)
 
-    # Connect buttons to functions
-    classify_btn.click(fn=predict, inputs=input_image_basic,
-                       outputs=output_label_basic, api_name="predict")
-    gradcam_btn.click(fn=generate_gradcam, inputs=input_image_gradcam, outputs=[
-                      output_gradcam, output_gradcam_info], api_name="gradcam")
-    shap_btn.click(fn=generate_shap, inputs=input_image_shap, outputs=[
-                   output_shap, output_shap_info], api_name="shap")
-    lime_btn.click(fn=generate_lime, inputs=input_image_lime, outputs=[
-                   output_lime_positive, output_lime_both, output_lime_info], api_name="lime")
+    # Connect button to unified function
+    predict_btn.click(
+        fn=predict_with_xai,
+        inputs=input_image,
+        outputs=[output_label, output_gradcam, output_shap,
+                 output_lime_positive, output_lime_both],
+        api_name="predict"
+    )
 
 # Launch with error reporting enabled
 if __name__ == "__main__":