JobenTan commited on
Commit
57bcc48
·
verified ·
1 Parent(s): d1a17ff

Upload 4 files

Browse files
Files changed (4) hide show
  1. app (1).py +76 -0
  2. gradcam.py +60 -0
  3. requirements (1).txt +5 -0
  4. vgg_xray_model.pth +3 -0
app (1).py ADDED
@@ -0,0 +1,76 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import torch.nn as nn
3
+ import torchvision.transforms as transforms
4
+ from torchvision.models import vgg16, VGG16_Weights
5
+ from PIL import Image
6
+ import gradio as gr
7
+ import numpy as np
8
+ import cv2
9
+ from gradcam import GradCAM, apply_heatmap
10
+
11
+ # Categories for classification
12
+ categories = ["COVID", "Lung_Opacity", "Normal", "Viral Pneumonia"]
13
+
14
+ # Load Pretrained VGG16 Model
15
+ print("Loading Model...")
16
+ device = torch.device("cpu")
17
+ vgg_model = vgg16(weights=VGG16_Weights.IMAGENET1K_V1)
18
+ for param in vgg_model.features.parameters():
19
+ param.requires_grad = False
20
+ vgg_model.classifier[6] = nn.Linear(4096, 4)
21
+ vgg_model.load_state_dict(torch.load("vgg_xray_model.pth", map_location=device))
22
+ vgg_model = vgg_model.to(device)
23
+ vgg_model.eval()
24
+ print("Model Loaded Successfully!")
25
+
26
+ # Initialize Grad-CAM with the last convolutional layer
27
+ gradcam = GradCAM(vgg_model, vgg_model.features[-1])
28
+
29
+ # Image Preprocessing
30
+ transform = transforms.Compose([
31
+ transforms.Resize((224, 224)),
32
+ transforms.ToTensor(),
33
+ transforms.Normalize(mean=[0.5, 0.5, 0.5], std=[0.5, 0.5, 0.5])
34
+ ])
35
+
36
+ # Prediction function with Grad-CAM
37
+ def predict(image):
38
+ # Preprocess image
39
+ image_tensor = transform(image).unsqueeze(0).to(device) # Add batch dimension
40
+
41
+ # Perform inference (no gradients)
42
+ with torch.no_grad():
43
+ output = vgg_model(image_tensor)
44
+ _, predicted = torch.max(output, 1)
45
+ prediction = categories[predicted.item()]
46
+
47
+ # Enable gradients for Grad-CAM
48
+ image_tensor.requires_grad_()
49
+ heatmap = gradcam.generate_heatmap(image_tensor, predicted.item())
50
+
51
+ # Overlay heatmap (handle failures)
52
+ try:
53
+ overlay = apply_heatmap(image, heatmap)
54
+ except:
55
+ overlay = image # Fallback to original image
56
+
57
+ return {"prediction": prediction}, overlay
58
+
59
+ # Gradio Interface
60
+ print("Initializing Gradio Interface...")
61
+
62
+ interface = gr.Interface(
63
+ fn=predict,
64
+ inputs=gr.Image(type="pil"),
65
+ outputs=[
66
+ gr.JSON(label="Diagnosis"), # Explicitly labeled JSON output
67
+ gr.Image(type="pil", label="Grad-CAM Heatmap")
68
+ ],
69
+ title="X-Ray Classifier with Grad-CAM",
70
+ allow_flagging="never", # Replace with flagging_mode="never" if using Gradio >=4.0
71
+ live=True
72
+ )
73
+
74
+ print("Launching Gradio App...")
75
+ interface.launch(server_name="0.0.0.0", server_port=7860)
76
+ print("Gradio App Running!")
gradcam.py ADDED
@@ -0,0 +1,60 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import torch
2
+ import numpy as np
3
+ import cv2
4
+ from PIL import Image
5
+
6
+ class GradCAM:
7
+ def __init__(self, model, target_layer):
8
+ self.model = model
9
+ self.target_layer = target_layer
10
+ self.gradients = None
11
+ self.activations = None
12
+
13
+ # Use register_full_backward_hook for PyTorch >=1.8
14
+ self.target_layer.register_forward_hook(self.save_activations)
15
+ self.target_layer.register_full_backward_hook(self.save_gradients)
16
+
17
+ def save_activations(self, module, input, output):
18
+ self.activations = output
19
+
20
+ def save_gradients(self, module, grad_input, grad_output):
21
+ self.gradients = grad_output[0] # grad_output is a tuple
22
+
23
+ def generate_heatmap(self, image_tensor, class_index=None):
24
+ # Ensure gradients are enabled
25
+ with torch.set_grad_enabled(True):
26
+ output = self.model(image_tensor) # No unsqueeze needed
27
+
28
+ if class_index is None:
29
+ class_index = output.argmax()
30
+
31
+ score = output[:, class_index]
32
+ self.model.zero_grad()
33
+ score.backward(retain_graph=True) # Compute gradients
34
+
35
+ # Move data to CPU and convert to numpy
36
+ gradients = self.gradients.cpu().numpy()
37
+ activations = self.activations.cpu().numpy()
38
+
39
+ # Compute weights and heatmap
40
+ weights = np.mean(gradients, axis=(2, 3)) # Shape: [batch, channels]
41
+ cam = np.sum(weights[:, :, None, None] * activations, axis=1) # Weighted sum
42
+ cam = np.maximum(cam, 0) # ReLU
43
+ cam = cam[0] # Remove batch dimension
44
+
45
+ # Normalize heatmap
46
+ cam = (cam - np.min(cam)) / (np.max(cam) - np.min(cam) + 1e-8)
47
+ return cam
48
+
49
+ def apply_heatmap(image_pil, heatmap):
50
+ """Overlay Grad-CAM heatmap on the original image"""
51
+ if heatmap is None: # Fallback if heatmap fails
52
+ return image_pil
53
+
54
+ image_cv = np.array(image_pil)
55
+ image_cv = cv2.cvtColor(image_cv, cv2.COLOR_RGB2BGR)
56
+ heatmap = cv2.resize(heatmap, (image_cv.shape[1], image_cv.shape[0]))
57
+ heatmap = np.uint8(255 * heatmap)
58
+ heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
59
+ overlay = cv2.addWeighted(image_cv, 0.5, heatmap, 0.5, 0)
60
+ return Image.fromarray(cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB))
requirements (1).txt ADDED
@@ -0,0 +1,5 @@
 
 
 
 
 
 
1
+ torch
2
+ torchvision
3
+ gradio==5.23.3
4
+ opencv-python-headless
5
+ anyio==3.6.2
vgg_xray_model.pth ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:492760a84b353e77027e0ddef767d976c240c97acbc41f3e88d6fa0267e85443
3
+ size 537119090