Spaces:
Build error
Build error
File size: 2,567 Bytes
6d08347 a76f433 d54549f a76f433 693b41f c00ea70 a76f433 c00ea70 a76f433 d9b8d56 a76f433 c00ea70 a76f433 c00ea70 6d08347 c00ea70 a76f433 d9b8d56 a76f433 c00ea70 a76f433 c00ea70 a76f433 c00ea70 a76f433 c00ea70 a76f433 c00ea70 a76f433 c00ea70 a76f433 c00ea70 a76f433 c00ea70 a76f433 c00ea70 a76f433 c00ea70 a76f433 c00ea70 a76f433 c00ea70 a76f433 c00ea70 a76f433 d516559 a76f433 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 |
import gradio as gr
import torch
import matplotlib.pyplot as plt
import numpy as np
from fastai.vision.all import load_learner, PILImage
import io
# Ensure custom classes exist before loading the model
class Hook:
def __init__(self, module, func):
self.hook = module.register_forward_hook(lambda mod, inp, out: func(out))
def remove(self): self.hook.remove()
class HookBwd:
def __init__(self, module, func):
self.hook = module.register_full_backward_hook(lambda mod, grad_input, grad_output: func(grad_output[0]))
def remove(self): self.hook.remove()
# Load the learner
try:
learn = load_learner('export.pkl')
print("Model loaded successfully!")
except Exception as e:
print(f"Error loading model: {e}")
# Function to predict + generate Class Activation Map (CAM)
def predict_with_cam(img):
img = PILImage.create(img)
# Get model and target layer (modify as needed for your architecture)
model = learn.model
target_layer = model[0][-1] # Adjust based on model architecture
activations, gradients = [], []
# Define hook functions
def hook_activations(out): activations.append(out)
def hook_gradients(grad): gradients.append(grad)
# Attach hooks
h1 = Hook(target_layer, hook_activations)
h2 = HookBwd(target_layer, hook_gradients)
# Run prediction
pred_class, pred_idx, probs = learn.predict(img)
# Perform backward pass for gradients
img_tensor = learn.dls.test_dl([img]).one_batch()[0]
img_tensor.requires_grad_()
output = model(img_tensor)
output[0, pred_idx].backward()
# Remove hooks
h1.remove()
h2.remove()
# Generate Class Activation Map (CAM)
act = activations[0].detach().cpu().squeeze(0)
grad = gradients[0].detach().cpu().squeeze(0)
weights = grad.mean(dim=(1, 2), keepdim=True)
cam = (weights * act).sum(0)
cam = cam.clamp(min=0).numpy()
# Normalize CAM
cam = (cam - cam.min()) / (cam.max() - cam.min())
# Plot CAM
fig, ax = plt.subplots()
ax.imshow(img)
ax.imshow(cam, alpha=0.5, cmap='jet')
ax.axis('off')
# Save CAM image
buf = io.BytesIO()
plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
buf.seek(0)
return {learn.dls.vocab[i]: float(probs[i]) for i in range(len(probs))}, buf
# Create 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.launch()
|