Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import torch
+import torch.nn.functional as F
 from torchvision import models, transforms
-from torch.nn import functional as F
 from PIL import Image
 import numpy as np
 import cv2
@@ -14,82 +14,79 @@ model.eval()
 
 classes = ['benign', 'malignant']
 
-# Preprocessing
+# Transform
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor()
 ])
 
-# Grad-CAM setup
-final_conv_layer = model.layer3[1].conv2  # Adjust if using a different architecture
-gradients = []
+# Target layer for Grad-CAM
+target_layer = model.layer3[1].conv2
 
-def save_gradient(module, grad_input, grad_output):
-    gradients.append(grad_output[0])
+# Store activations & gradients
+activations = None
+gradients = None
 
-final_conv_layer.register_forward_hook(save_gradient)
+def forward_hook(module, input, output):
+    global activations
+    activations = output.detach()
 
-def generate_gradcam(input_tensor, pred_class):
+def backward_hook(module, grad_input, grad_output):
+    global gradients
+    gradients = grad_output[0].detach()
+
+target_layer.register_forward_hook(forward_hook)
+target_layer.register_backward_hook(backward_hook)
+
+# Grad-CAM function
+def generate_gradcam(input_tensor, class_idx):
     model.zero_grad()
     output = model(input_tensor)
-    class_score = output[0, pred_class]
+    class_score = output[0, class_idx]
     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]
-
+    pooled_gradients = torch.mean(gradients, dim=[0, 2, 3])  # [C]
+    weighted_activations = activations[0] * pooled_gradients[:, None, None]  # [C, H, W]
+    cam = torch.sum(weighted_activations, dim=0).cpu().numpy()
+    
+    # Normalize and resize
     cam = np.maximum(cam, 0)
     cam = cv2.resize(cam, (224, 224))
-    cam = cam - np.min(cam)
-    cam = cam / np.max(cam)
-
+    cam -= cam.min()
+    cam /= cam.max()
     return cam
 
-# Hook to get activations
-def get_activations(module, input, output):
-    global final_conv_layer_output
-    final_conv_layer_output = output
+# Full pipeline
+def predict(img):
+    global activations, gradients
+    activations = None
+    gradients = None
 
-final_conv_layer.register_forward_hook(get_activations)
+    img = img.convert("RGB")
+    input_tensor = transform(img).unsqueeze(0)
 
-# Main prediction and Grad-CAM overlay function
-def predict(img):
-    gradients.clear()
-    img_rgb = img.convert("RGB")
-    input_tensor = transform(img_rgb).unsqueeze(0)
-    
-    with torch.no_grad():
-        output = model(input_tensor)
-        probs = F.softmax(output[0], dim=0)
-        pred_class = torch.argmax(probs).item()
+    output = model(input_tensor)
+    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)
+    # Convert to heatmap
+    heatmap = np.uint8(255 * cam)
     heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
 
-    # Overlay heatmap on image
-    img_np = np.array(img_rgb.resize((224, 224)))
+    # Overlay
+    img_np = np.array(img.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
-title = "🧠 Lumen: Skin Cancer Classifier with Grad-CAM"
+# Gradio interface
+title = "🧠 Lumen: Skin Cancer Classifier + Grad-CAM"
 description = """
 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.
+<b>Disclaimer:</b> For educational use only.
 """
 
 demo = gr.Interface(