Update app.py

app.py

@@ -139,20 +139,23 @@
 
 import gradio as gr
 import numpy as np
-from PIL import Image
-import cv2
 import tensorflow as tf
 from tensorflow.keras.models import load_model
-from gradcam_utils import generate_heatmap_tf_explain
+from PIL import Image
+import os
+
 from models import create_vgg19_model
+from gradcam_utils import generate_heatmap_tf_explain
 
 # Load your trained model
 ensemble_model = load_model("ensemble_model_best(92.3).h5")
-vgg_model = create_vgg19_model()
+vgg_model = create_vgg19_model()  # Only used for Grad-CAM (tf-explain)
 
+# Label names
 def get_class_name(class_id):
     return "Normal" if class_id == 0 else "Pneumonia"
 
+# Prediction + Heatmap generation
 def predict_and_heatmap(image):
     img = image.resize((224, 224))
     img_array = np.array(img) / 255.0
@@ -162,7 +165,6 @@ def predict_and_heatmap(image):
     class_id = int(np.argmax(prediction[0]))
     label = get_class_name(class_id)
 
-    # Styled HTML result
     result_html = f"""
     <div style='
         text-align: center;
@@ -180,11 +182,15 @@ def predict_and_heatmap(image):
     </div>
     """
 
-    # Generate heatmap from VGG19 using tf-explain
+    # Generate Grad-CAM heatmap using tf-explain (on VGG19)
     heatmap_img = generate_heatmap_tf_explain(image, vgg_model, class_index=class_id)
     return result_html, heatmap_img
 
-# Styled Gradio Interface
+# Function to load sample image
+def load_sample():
+    return Image.open("sample_pneumonia.jpeg")
+
+# Gradio interface
 with gr.Blocks(theme="soft") as demo:
     gr.Markdown("""
     <div style="text-align: center; font-size: 2.5rem; font-weight: bold; color: #0b5394; margin-bottom: 1rem;">
@@ -199,12 +205,16 @@ with gr.Blocks(theme="soft") as demo:
         with gr.Column(scale=1, min_width=600):
             image_input = gr.Image(type="pil", label="Upload Chest X-Ray", interactive=True, width=600, height=600)
             prediction_output = gr.HTML(label="Prediction")
-            heatmap_output = gr.Image(label="Grad-CAM Heatmap")
-            submit_button = gr.Button("Predict")
-            clear_button = gr.Button("Clear")
+            heatmap_output = gr.Image(label="Grad-CAM Heatmap", width=600, height=600)
+
+            with gr.Row():
+                submit_button = gr.Button("Predict")
+                clear_button = gr.Button("Clear")
+                sample_button = gr.Button("Load Sample X-ray")
 
     submit_button.click(fn=predict_and_heatmap, inputs=image_input, outputs=[prediction_output, heatmap_output])
     clear_button.click(fn=lambda: (None, "", None), inputs=[], outputs=[image_input, prediction_output, heatmap_output])
+    sample_button.click(fn=load_sample, inputs=[], outputs=[image_input])
 
     gr.Markdown("""
     <div style="text-align: center; font-size: 0.95rem; color: #888; margin-top: 30px;">