Update app.py

app.py

@@ -91,6 +91,44 @@ def superimpose_images(base_image, overlay_image, alpha):
 
     return Image.fromarray(blended_array)
 
+def compute_gradcam(image_tensor, model):
+    target_layer = model.layer4[-1]
+    gradients, activations = [], []
+
+    def backward_hook(module, grad_input, grad_output):
+        gradients.append(grad_output[0])
+
+    def forward_hook(module, input, output):
+        activations.append(output)
+
+    target_layer.register_forward_hook(forward_hook)
+    target_layer.register_backward_hook(backward_hook)
+    
+    output = model(image_tensor)
+    class_idx = output.argmax(dim=1).item()
+    model.zero_grad()
+    output[:, class_idx].backward()
+
+    grads = gradients[0].cpu().data.numpy()[0]
+    acts = activations[0].cpu().data.numpy()[0]
+    weights = np.mean(grads, axis=(1, 2))
+    cam = np.zeros(acts.shape[1:], dtype=np.float32)
+
+    for i, w in enumerate(weights):
+        cam += w * acts[i]
+    
+    cam = np.maximum(cam, 0)
+    cam = cv2.resize(cam, (224, 224))
+    cam = (cam - cam.min()) / (cam.max() - cam.min())
+    return cam
+
+def overlay_heatmap(image, heatmap):
+    heatmap = cv2.applyColorMap(np.uint8(255 * heatmap), cv2.COLORMAP_JET)
+    heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
+    image = np.array(image.resize((224, 224)))
+    superimposed_img = cv2.addWeighted(image, 0.6, heatmap, 0.4, 0)
+    return Image.fromarray(superimposed_img)
+
 # Prediction function
 def predict(image, brightness, contrast, hue, overlay_image, alpha):
     """Apply filters, superimpose, classify image, and visualize results."""
@@ -110,6 +148,9 @@ def predict(image, brightness, contrast, hue, overlay_image, alpha):
         output = model(image_tensor)
         probabilities = F.softmax(output, dim=1).cpu().numpy()[0]
 
+    heatmap = compute_gradcam(image_tensor, model)
+    heatmap_image = overlay_heatmap(final_image, heatmap)
+
     # Generate Bar Chart
     with plt.xkcd():
         fig, ax = plt.subplots(figsize=(5, 3))
@@ -120,7 +161,7 @@ def predict(image, brightness, contrast, hue, overlay_image, alpha):
         for i, v in enumerate(probabilities):
             ax.text(i, v + 0.02, f"{v:.2f}", ha='center', fontsize=10)
 
-    return final_image, fig
+    return final_image, heatmap_image, fig
 
 # Gradio Interface
 with gr.Blocks() as interface:
@@ -137,6 +178,7 @@ with gr.Blocks() as interface:
 
         with gr.Column():
             processed_image = gr.Image(label="Final Processed Image")
+            heatmap_image = gr.Image(label="Heatmap Visualization")
             bar_chart = gr.Plot(label="Class Probabilities")
 
     inputs = [image_input, brightness, contrast, hue, overlay_input, alpha]