Update app.py

app.py

@@ -0,0 +1,88 @@
+import gradio as gr
+import torch
+import torch.nn as nn
+import numpy as np
+from torchvision import transforms
+from transformers import ViTForImageClassification
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.image import show_cam_on_image
+from PIL import Image
+
+# Device
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# Transforms
+test_transforms = transforms.Compose([
+    transforms.Resize((224, 224)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                         std=[0.229, 0.224, 0.225])
+])
+
+mean = np.array([0.485, 0.456, 0.406])
+std  = np.array([0.229, 0.224, 0.225])
+
+# Load model
+model = ViTForImageClassification.from_pretrained(
+    'google/vit-base-patch16-224',
+    num_labels=2,
+    ignore_mismatched_sizes=True
+)
+model.load_state_dict(torch.load('vit_pneumonia_best.pth', map_location=device))
+model = model.to(device)
+model.eval()
+
+# ViT wrapper for Grad-CAM
+class ViTWrapper(nn.Module):
+    def __init__(self, model):
+        super().__init__()
+        self.model = model
+    def forward(self, x):
+        return self.model(x).logits
+
+def reshape_transform(tensor, height=14, width=14):
+    result = tensor[:, 1:, :].reshape(tensor.size(0), height, width, tensor.size(2))
+    result = result.transpose(2, 3).transpose(1, 2)
+    return result
+
+wrapped_model = ViTWrapper(model)
+target_layers = [model.vit.encoder.layer[-1].layernorm_before]
+cam = GradCAM(model=wrapped_model, target_layers=target_layers,
+              reshape_transform=reshape_transform)
+
+# Prediction function
+def predict(pil_image):
+    img_tensor = test_transforms(pil_image).unsqueeze(0).to(device)
+
+    with torch.no_grad():
+        output = model(img_tensor).logits
+        probs  = torch.softmax(output, dim=1)[0]
+
+    grayscale_cam = cam(input_tensor=img_tensor)[0]
+
+    img_numpy = test_transforms(pil_image).permute(1, 2, 0).numpy()
+    img_numpy = std * img_numpy + mean
+    img_numpy = np.clip(img_numpy, 0, 1).astype(np.float32)
+
+    heatmap = show_cam_on_image(img_numpy, grayscale_cam, use_rgb=True)
+
+    label = {
+        "NORMAL":    float(probs[0]),
+        "PNEUMONIA": float(probs[1])
+    }
+
+    return label, Image.fromarray(heatmap)
+
+# Gradio interface
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil", label="Upload Chest X-Ray"),
+    outputs=[
+        gr.Label(label="Prediction"),
+        gr.Image(type="pil", label="Grad-CAM Heatmap")
+    ],
+    title="🫁 Pneumonia Detector",
+    description="Upload a chest X-ray image to detect pneumonia with explainable AI (Grad-CAM)",
+)
+
+demo.launch()