Update app.py

app.py

@@ -3,6 +3,10 @@ import gradio as gr
 import timm
 from torchvision import transforms as T
 from PIL import Image
+import numpy as np
+import cv2
+from pytorch_grad_cam import GradCAMPlusPlus
+from pytorch_grad_cam.utils.image import show_cam_on_image
 
 # ==========================================
 # 1. Configuration & Setup
@@ -42,49 +46,45 @@ transform = T.Compose([
 ])
 
 # ==========================================
-# 4. Styled Text Blocks & Inference
+# 4. Inference & Grad-CAM Function
 # ==========================================
-# Styled "Waiting" Box - Soft light gray background with dashed border
-waiting_text = """
-<div style="background-color: #f1f5f9; padding: 15px; border-radius: 10px; border: 2px dashed #cbd5e1; text-align: center; margin-bottom: 15px;">
-    <h3 style="color: #64748b; margin: 0; font-weight: 600;">⏳ Waiting for upload...</h3>
-</div>
-"""
+waiting_text = "<div style='background: #f8fafc; color: #475569; padding: 10px; border-radius: 5px;'><h3>Waiting for upload...</h3></div>"
 
 def predict(image):
     if image is None:
-        return None, waiting_text
+        return None, waiting_text, None
 
+    # Prepare Image
     image = image.convert("RGB")
     input_tensor = transform(image).unsqueeze(0).to(device)
 
+    # 1. Classification (Inside no_grad)
     with torch.no_grad():
         outputs = model(input_tensor)
         probabilities = torch.nn.functional.softmax(outputs[0], dim=0)
 
-    # Format for Gradio Label component
-    confidences = {class_names[i]: float(
-        probabilities[i]) for i in range(len(class_names))}
-
-    # Get top prediction for the text output
+    confidences = {class_names[i]: float(probabilities[i]) for i in range(len(class_names))}
     top_class = max(confidences, key=confidences.get)
     top_score = confidences[top_class]
     
-    # Styled Success Box - Soft light indigo background, dark indigo text, nice shadow
-    message = f"""
-    <div style="background-color: #e0e7ff; padding: 15px; border-radius: 10px; border-left: 6px solid #4f46e5; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1); text-align: center; margin-bottom: 15px;">
-        <h3 style="color: #312e81; margin: 0; font-weight: 800; font-size: 1.2rem;">
-            🧠 The model is {top_score*100:.1f}% confident that the primary expression is {top_class}.
-        </h3>
-    </div>
-    """
+    message = f"<div style='background: #e0e7ff; color: #1e3a8a; padding: 10px; border-radius: 5px;'><h3>🧠 The model is <b>{top_score*100:.1f}%</b> confident that the primary expression is <b>{top_class}</b>.</h3></div>"
 
-    return confidences, message
+    # 2. Grad-CAM Generation (Outside no_grad because Grad-CAM requires gradients)
+    target_layers = [model.layer4[-1].conv3]
+    cam = GradCAMPlusPlus(model=model, target_layers=target_layers, use_cuda=False) # CPU mode
+    
+    # Generate heatmap
+    grayscale_cam = cam(input_tensor=input_tensor)[0, :]
+    
+    # Overlay heatmap on original image
+    rgb_img = np.array(image.resize((im_size, im_size))) / 255.0
+    cam_image = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True, image_weight=0.4)
+
+    return confidences, message, cam_image
 
 # ==========================================
 # 5. Modern Gradio UI (Blocks)
 # ==========================================
-# Using a sleek, soft theme
 theme = gr.themes.Soft(
     primary_hue="indigo",
     secondary_hue="blue",
@@ -93,32 +93,14 @@ theme = gr.themes.Soft(
 )
 
 custom_css = """
-.gradio-container {
-    background-color: #f8fafc;
-}
-.header-box {
-    text-align: center;
-    padding: 2rem;
-    background: linear-gradient(135deg, #4f46e5 0%, #3b82f6 100%);
-    border-radius: 15px;
-    color: white;
-    box-shadow: 0 10px 15px -3px rgba(0, 0, 0, 0.1);
-    margin-bottom: 20px;
-}
-.header-title {
-    font-size: 2.5rem;
-    font-weight: 800;
-    margin-bottom: 0.5rem;
-}
-.header-subtitle {
-    font-size: 1.1rem;
-    opacity: 0.9;
-}
+.gradio-container { background-color: #f8fafc; }
+.header-box { text-align: center; padding: 2rem; background: linear-gradient(135deg, #4f46e5 0%, #3b82f6 100%); border-radius: 15px; color: white; box-shadow: 0 10px 15px -3px rgba(0, 0, 0, 0.1); margin-bottom: 20px; }
+.header-title { font-size: 2.5rem; font-weight: 800; margin-bottom: 0.5rem; }
+.header-subtitle { font-size: 1.1rem; opacity: 0.9; }
 """
 
 with gr.Blocks(theme=theme, css=custom_css) as demo:
 
-    # --- Custom Header ---
     gr.HTML("""
     <div class="header-box">
         <div class="header-title">✨ Autistic Children Expression Recognition</div>
@@ -126,59 +108,47 @@ with gr.Blocks(theme=theme, css=custom_css) as demo:
     </div>
     """)
 
-    # --- Main App Layout ---
     with gr.Row():
-
         # Left Column: Inputs
         with gr.Column(scale=1):
-            # Styled Upload Header - White card, subtle shadow, indigo accent
-            gr.HTML("""
-            <div style="background-color: #ffffff; padding: 15px; border-radius: 10px; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1); border-left: 5px solid #4f46e5; margin-bottom: 15px;">
-                <h3 style="color: #1e293b; margin: 0 0 5px 0; font-weight: 800;">📸 Upload Image</h3>
-                <p style="color: #475569; margin: 0; font-size: 0.95rem;">Upload a clear photo of a child's face to begin the analysis.</p>
-            </div>
-            """)
-            input_image = gr.Image(
-                type="pil", label="Input Image", elem_classes="image-box")
+            gr.HTML("<div style='background: #f1f5f9; color: #0f172a; padding: 10px; border-radius: 5px;'><h3>📸 Upload Image</h3>Upload a clear photo of a child's face.</div>")
+            input_image = gr.Image(type="pil", label="Input Image", elem_classes="image-box")
 
             with gr.Row():
                 clear_btn = gr.Button("🗑️ Clear", variant="secondary")
-                submit_btn = gr.Button(
-                    "🚀 Analyze Expression", variant="primary")
+                submit_btn = gr.Button("🚀 Analyze Expression", variant="primary")
 
-        # Right Column: Outputs
+        # Right Column: Outputs (Parallel to Input)
         with gr.Column(scale=1):
-            # Styled Results Header - White card, subtle shadow, blue accent
-            gr.HTML("""
-            <div style="background-color: #ffffff; padding: 15px; border-radius: 10px; box-shadow: 0 4px 6px -1px rgba(0, 0, 0, 0.1); border-left: 5px solid #3b82f6; margin-bottom: 15px;">
-                <h3 style="color: #1e293b; margin: 0; font-weight: 800;">📊 AI Analysis Results</h3>
-            </div>
-            """)
+            gr.HTML("<div style='background: #f1f5f9; color: #0f172a; padding: 10px; border-radius: 5px;'><h3>📊 AI Analysis & Grad-CAM</h3></div>")
             
+            # Message and Label block
             output_message = gr.HTML(waiting_text)
-            output_label = gr.Label(
-                num_top_classes=4, label="Confidence Scores")
+            
+            with gr.Row():
+                output_label = gr.Label(num_top_classes=4, label="Confidence Scores")
+                output_cam = gr.Image(type="numpy", label="Grad-CAM Focus Map")
 
-    # --- Expandable Info Section ---
     with gr.Accordion("ℹ️ About this Model", open=False):
         gr.Markdown("""
         * **Architecture:** ResNeXt-101 (32x8d)
         * **Input Resolution:** 224x224 pixels
+        * **Visualizer:** Grad-CAM++ (Highlights where the AI is looking)
         * **Task:** Multi-class Facial Expression Classification
         * **Note:** This tool is for educational/research demonstration purposes and is not a clinical diagnostic tool.
         """)
 
-    # --- Button Functionality ---
+    # Button Functionality
     submit_btn.click(
         fn=predict,
         inputs=[input_image],
-        outputs=[output_label, output_message]
+        outputs=[output_label, output_message, output_cam]
     )
 
     clear_btn.click(
-        fn=lambda: (None, None, waiting_text),
+        fn=lambda: (None, waiting_text, None),
         inputs=[],
-        outputs=[input_image, output_label, output_message]
+        outputs=[input_image, output_label, output_message, output_cam]
     )
 
 if __name__ == "__main__":