Update app.py

app.py

@@ -9,6 +9,8 @@ from albumentations.pytorch import ToTensorV2
 from torchvision import transforms
 import segmentation_models_pytorch as smp
 import torch.nn as nn
+import cv2
+import matplotlib.cm as cm
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -41,13 +43,52 @@ classifier_transform = transforms.Compose([
     transforms.Normalize(mean=0.41, std=0.16)
 ])
 
+# Grad-CAM function
+def generate_gradcam(model, image_tensor, target_class):
+    gradients = []
+    activations = []
+
+    def backward_hook(module, grad_input, grad_output):
+        gradients.append(grad_output[0])
+
+    def forward_hook(module, input, output):
+        activations.append(output)
+
+    final_conv = model.features[-1]
+    handle_fwd = final_conv.register_forward_hook(forward_hook)
+    handle_bwd = final_conv.register_backward_hook(backward_hook)
+
+    model.zero_grad()
+    output = model(image_tensor)
+    class_score = output[0, target_class]
+    class_score.backward()
+
+    grads_val = gradients[0].cpu().numpy()[0]
+    activations_val = activations[0].cpu().numpy()[0]
+
+    weights = np.mean(grads_val, axis=(1, 2))
+    cam = np.zeros(activations_val.shape[1:], dtype=np.float32)
+    for i, w in enumerate(weights):
+        cam += w * activations_val[i]
+
+    cam = np.maximum(cam, 0)
+    cam = cv2.resize(cam, (224, 224))
+    cam = cam - np.min(cam)
+    cam = cam / np.max(cam + 1e-8)
+
+    heatmap = cm.jet(cam)[:, :, :3]  # RGB heatmap
+    heatmap = (heatmap * 255).astype(np.uint8)
+
+    handle_fwd.remove()
+    handle_bwd.remove()
+
+    return heatmap
+
 # Inference Function
 def analyze(image):
-    # Convert to grayscale for UNet
     image_gray = image.convert("L")
     img_np = np.array(image_gray)
 
-    # Preprocess for UNet
     augmented = unet_transform(image=img_np)
     unet_input = augmented["image"].unsqueeze(0).to(device)
 
@@ -56,37 +97,39 @@ def analyze(image):
         mask_pred = torch.sigmoid(mask_pred)
         mask_binary = (mask_pred > 0.5).float()
 
-    # Resize original image and mask to 224x224 for classifier
     resized = image_gray.resize((224, 224), Image.BILINEAR)
     mask_resized = transforms.functional.resize(mask_binary, [224, 224])
-    
-    image_np = np.array(resized).astype(np.float32)  # shape: (224, 224)
-    mask_np = (mask_resized.squeeze().cpu().numpy() > 0.5).astype(np.float32)  # shape: (224, 224)
 
-    # Apply mask and convert to 3-channel image
+    image_np = np.array(resized).astype(np.float32)
+    mask_np = (mask_resized.squeeze().cpu().numpy() > 0.5).astype(np.float32)
+
     lung_image = image_np * mask_np
-    lung_image_3ch = np.stack([lung_image] * 3, axis=-1)  # shape: (224, 224, 3)
+    lung_image_3ch = np.stack([lung_image] * 3, axis=-1)
     lung_image_3ch = np.clip(lung_image_3ch, 0, 255).astype(np.uint8)
     lung_image_pil = Image.fromarray(lung_image_3ch)
 
-    # Classification
     input_tensor = classifier_transform(lung_image_pil).unsqueeze(0).to(device)
 
-    with torch.no_grad():
-        logits = m2(input_tensor)
-        probs = torch.softmax(logits, dim=1)
-        confidence, pred_class = torch.max(probs, dim=1)
+    logits = m2(input_tensor)
+    probs = torch.softmax(logits, dim=1)
+    confidence, pred_class = torch.max(probs, dim=1)
+
+    heatmap = generate_gradcam(m2, input_tensor, pred_class.item())
+    overlay = cv2.addWeighted(lung_image_3ch, 0.5, heatmap, 0.5, 0)
+    overlay_pil = Image.fromarray(overlay)
 
     classes = ["COVID", "Lung_Opacity", "Normal", "Viral Pneumonia"]
-    return f"Prediction: {classes[pred_class.item()]} ({confidence.item()*100:.2f}%)"
+    result_text = f"Prediction: {classes[pred_class.item()]} ({confidence.item()*100:.2f}%)"
+
+    return result_text, overlay_pil
 
 # Gradio UI
 interface = gr.Interface(
     fn=analyze,
     inputs=gr.Image(type="pil"),
-    outputs="text",
-    title="Chest X-Ray Analysis",
-    description="Upload a chest X-ray to detect disease based on segmentation and classification models."
+    outputs=[gr.Text(), gr.Image(type="pil")],
+    title="Chest X-Ray Analysis with Grad-CAM",
+    description="Upload a chest X-ray to detect disease and view model attention using Grad-CAM."
 )
 
-interface.launch()
+interface.launch()