Update app.py

app.py

@@ -1,7 +1,10 @@
 import gradio as gr
 import torch
 from torchvision import models, transforms
+from torch.nn import functional as F
 from PIL import Image
+import numpy as np
+import cv2
 
 # Load model
 model = models.resnet18()
@@ -9,37 +12,93 @@ model.fc = torch.nn.Linear(model.fc.in_features, 2)
 model.load_state_dict(torch.load("skin_cancer_resnet18_version1.pt", map_location="cpu"))
 model.eval()
 
-# Class labels
 classes = ['benign', 'malignant']
 
-# Image preprocessing
+# Preprocessing
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor()
 ])
 
-# Inference function
+# Grad-CAM setup
+final_conv_layer = model.layer4[1].conv2  # Adjust if using a different architecture
+gradients = []
+
+def save_gradient(module, grad_input, grad_output):
+    gradients.append(grad_output[0])
+
+final_conv_layer.register_forward_hook(save_gradient)
+
+def generate_gradcam(input_tensor, pred_class):
+    model.zero_grad()
+    output = model(input_tensor)
+    class_score = output[0, pred_class]
+    class_score.backward()
+
+    grads_val = gradients[-1].detach().numpy()[0]
+    activations = final_conv_layer_output.detach().numpy()[0]
+
+    weights = np.mean(grads_val, axis=(1, 2))
+    cam = np.zeros(activations.shape[1:], dtype=np.float32)
+
+    for i, w in enumerate(weights):
+        cam += w * activations[i]
+
+    cam = np.maximum(cam, 0)
+    cam = cv2.resize(cam, (224, 224))
+    cam = cam - np.min(cam)
+    cam = cam / np.max(cam)
+
+    return cam
+
+# Hook to get activations
+def get_activations(module, input, output):
+    global final_conv_layer_output
+    final_conv_layer_output = output
+
+final_conv_layer.register_forward_hook(get_activations)
+
+# Main prediction and Grad-CAM overlay function
 def predict(img):
-    img = img.convert("RGB")
-    input_tensor = transform(img).unsqueeze(0)
+    gradients.clear()
+    img_rgb = img.convert("RGB")
+    input_tensor = transform(img_rgb).unsqueeze(0)
+    
     with torch.no_grad():
         output = model(input_tensor)
-        probs = torch.nn.functional.softmax(output[0], dim=0)
-    return {classes[i]: float(probs[i]) for i in range(2)}
+        probs = F.softmax(output[0], dim=0)
+        pred_class = torch.argmax(probs).item()
+
+    # Grad-CAM
+    cam = generate_gradcam(input_tensor, pred_class)
+
+    # Convert CAM to heatmap
+    heatmap = (cam * 255).astype(np.uint8)
+    heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
+
+    # Overlay heatmap on image
+    img_np = np.array(img_rgb.resize((224, 224)))
+    overlay = cv2.addWeighted(img_np, 0.6, heatmap, 0.4, 0)
+
+    # Convert back to PIL
+    overlay_img = Image.fromarray(cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB))
+
+    return {classes[i]: float(probs[i]) for i in range(2)}, overlay_img
 
-# UI text
-title = "🧠 Lumen: Skin Cancer Classifier"
+# UI
+title = "🧠 Lumen: Skin Cancer Classifier with Grad-CAM"
 description = """
-Upload a dermoscopic image of a mole or skin lesion.<br>
-The model will classify it as <b>benign</b> or <b>malignant</b> based on its appearance.<br><br>
-<b>Disclaimer:</b> This tool is for research and educational use only. It is not a diagnostic device.
+Upload a dermoscopic image of a mole or lesion. The model will classify it as <b>benign</b> or <b>malignant</b> and show a heatmap of what it focused on.<br><br>
+<b>Disclaimer:</b> This tool is for educational use only.
 """
 
-# Gradio Interface
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil", label="Upload Lesion Image"),
-    outputs=gr.Label(num_top_classes=2, label="Prediction"),
+    outputs=[
+        gr.Label(num_top_classes=2, label="Prediction"),
+        gr.Image(type="pil", label="Grad-CAM Visualisation")
+    ],
     title=title,
     description=description
 )