Update app.py

app.py

@@ -88,29 +88,44 @@ transform = transforms.Compose([
 def predict_with_gradcam(image):
     image = image.convert("RGB")
     input_tensor = transform(image).unsqueeze(0).to(device)
+
+    # Prediction
     output = model(input_tensor)
     pred_idx = output.argmax(dim=1).item()
     pred_label = class_names[pred_idx]
 
-    gradcam = GradCAM(model, model.conv4)
-    cam = gradcam.generate(input_tensor)
-    cam_resized = cv2.resize(cam, (224, 224))
-    
+    # Prepare base image
     img_np = np.array(image.resize((224, 224))) / 255.0
-    cam_overlay = cv2.applyColorMap(np.uint8(255 * cam_resized), cv2.COLORMAP_JET)
-    cam_overlay = cv2.cvtColor(cam_overlay, cv2.COLOR_BGR2RGB)
-    overlay = (0.5 * img_np + 0.5 * cam_overlay / 255.0)
-    overlay = np.clip(overlay, 0, 1)
 
-    return pred_label, overlay
+    # Multiple layer Grad-CAMs
+    target_layers = [model.conv2, model.conv3, model.conv4]
+    visualizations = []
+
+    for layer in target_layers:
+        gradcam = GradCAM(model, layer)
+        cam = gradcam.generate(input_tensor)
+        cam_resized = cv2.resize(cam, (224, 224))
+        cam_overlay = cv2.applyColorMap(np.uint8(255 * cam_resized), cv2.COLORMAP_JET)
+        cam_overlay = cv2.cvtColor(cam_overlay, cv2.COLOR_BGR2RGB)
+        overlay = (0.5 * img_np + 0.5 * cam_overlay / 255.0)
+        overlay = np.clip(overlay, 0, 1)
+        visualizations.append(overlay)
+
+    return pred_label, *visualizations
+
 
-# Gradio interface
 interface = gr.Interface(
     fn=predict_with_gradcam,
     inputs=gr.Image(type="pil"),
-    outputs=["label", "image"],
+    outputs=[
+        gr.Label(label="Predicted Class"),
+        gr.Image(label="Grad-CAM: Conv2"),
+        gr.Image(label="Grad-CAM: Conv3"),
+        gr.Image(label="Grad-CAM: Conv4")
+    ],
     title="🦷 Teeth Disease Classifier with Grad-CAM",
-    description="Upload an image of teeth and the model will predict the disease with Grad-CAM visualization."
+    description="Upload a teeth image. The model predicts the class and shows Grad-CAM visualizations for multiple convolutional layers."
 )
 
+
 interface.launch()