Update app.py

app.py

@@ -1,17 +1,89 @@
 import gradio as gr
+import torch
+import matplotlib.pyplot as plt
+import numpy as np
 from fastai.vision.all import load_learner, PILImage
+from fastai.vision.utils import show_image
+import io
+from torchvision.transforms.functional import to_pil_image
 
-# Import any custom code here (if needed)
-# Example:
-# from your_module import CustomTransform, CustomModel
+# Hook classes from your notebook
+class Hook:
+    def __init__(self, m, f):
+        self.hook = m.register_forward_hook(lambda m, i, o: f(o))
+    def remove(self): self.hook.remove()
 
+class HookBwd:
+    def __init__(self, m, f):
+        self.hook = m.register_backward_hook(lambda m, gi, go: f(go[0]))
+    def remove(self): self.hook.remove()
+
+# Load the learner
 learn = load_learner('export.pkl')
 
-def predict(img):
+# Function to predict + generate CAM
+def predict_with_cam(img):
     img = PILImage.create(img)
-    pred, pred_idx, probs = learn.predict(img)
-    labels = learn.dls.vocab
-    return {labels[i]: float(probs[i]) for i in range(len(labels))}
+    
+    # Get the model and target layer (adjust depending on your model)
+    model = learn.model
+    target_layer = model[0][-1]  # Might need to adjust based on your architecture
+    
+    # Placeholders for activations and gradients
+    activations = []
+    gradients = []
+    
+    # Hook functions
+    def hook_activations(out): activations.append(out)
+    def hook_gradients(grad): gradients.append(grad)
+    
+    # Register hooks
+    h1 = Hook(target_layer, hook_activations)
+    h2 = HookBwd(target_layer, hook_gradients)
+    
+    # Prediction
+    pred_class, pred_idx, probs = learn.predict(img)
+    
+    # Backward pass to get gradients
+    output = learn.model(img.unsqueeze(0)) if not isinstance(img, torch.Tensor) else learn.model(img)
+    output[0, pred_idx].backward()
+    
+    # Remove hooks
+    h1.remove()
+    h2.remove()
+    
+    # Generate CAM
+    act = activations[0].detach().cpu()[0]
+    grad = gradients[0].detach().cpu()[0]
+    
+    weights = grad.mean(dim=(1, 2), keepdim=True)
+    cam = (weights * act).sum(0)
+    cam = cam.clamp(min=0).numpy()
+    
+    # Normalize CAM
+    cam -= cam.min()
+    cam /= cam.max()
+    
+    # Convert to image
+    fig, ax = plt.subplots()
+    ax.imshow(img)
+    ax.imshow(cam, alpha=0.5, cmap='jet')
+    ax.axis('off')
+    
+    # Save CAM to an in-memory buffer
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
+    
+    # Return predictions + CAM image
+    return {learn.dls.vocab[i]: float(probs[i]) for i in range(len(probs))}, buf
+
+# Gradio interface
+interface = gr.Interface(
+    fn=predict_with_cam,
+    inputs=gr.Image(type='pil'),
+    outputs=[gr.Label(num_top_classes=3), gr.Image(type='pil')],
+    title="Image Classifier with CAM"
+)
 
-interface = gr.Interface(fn=predict, inputs=gr.Image(), outputs=gr.Label(num_top_classes=3))
-interface.launch(share=True)
+interface.launch()