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()