File size: 4,860 Bytes
147fbb0 19f8d94 147fbb0 19f8d94 147fbb0 19f8d94 147fbb0 19f8d94 147fbb0 19f8d94 147fbb0 19f8d94 147fbb0 19f8d94 |
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 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 |
import gradio as gr
import timm
import torch
from PIL import Image
import requests
from io import BytesIO
import numpy as np
from pytorch_grad_cam import GradCAM, HiResCAM, ScoreCAM, GradCAMPlusPlus, AblationCAM, XGradCAM, EigenCAM, FullGrad
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
from pytorch_grad_cam.utils.image import show_cam_on_image
from timm.data import create_transform
# List of available timm models
MODELS = timm.list_models()
# List of available GradCAM methods
CAM_METHODS = {
"GradCAM": GradCAM,
"HiResCAM": HiResCAM,
"ScoreCAM": ScoreCAM,
"GradCAM++": GradCAMPlusPlus,
"AblationCAM": AblationCAM,
"XGradCAM": XGradCAM,
"EigenCAM": EigenCAM,
"FullGrad": FullGrad
}
def load_model(model_name):
model = timm.create_model(model_name, pretrained=True)
model.eval()
return model
def process_image(image_path, model):
if image_path.startswith('http'):
response = requests.get(image_path)
image = Image.open(BytesIO(response.content))
else:
image = Image.open(image_path)
config = model.pretrained_cfg
transform = create_transform(
input_size=config['input_size'],
crop_pct=config['crop_pct'],
mean=config['mean'],
std=config['std'],
interpolation=config['interpolation'],
is_training=False
)
tensor = transform(image).unsqueeze(0)
return tensor
def get_cam_image(model, image, target_layer, cam_method):
cam = CAM_METHODS[cam_method](model=model, target_layers=[target_layer])
grayscale_cam = cam(input_tensor=image)
config = model.pretrained_cfg
mean = torch.tensor(config['mean']).view(3, 1, 1)
std = torch.tensor(config['std']).view(3, 1, 1)
rgb_img = (image.squeeze(0) * std + mean).permute(1, 2, 0).cpu().numpy()
rgb_img = np.clip(rgb_img, 0, 1)
cam_image = show_cam_on_image(rgb_img, grayscale_cam[0, :], use_rgb=True)
return Image.fromarray(cam_image)
def get_feature_info(model):
if hasattr(model, 'feature_info'):
return [f['module'] for f in model.feature_info]
else:
return []
def get_target_layer(model, target_layer_name):
if target_layer_name is None:
return None
try:
return model.get_submodule(target_layer_name)
except AttributeError:
print(f"WARNING: Layer '{target_layer_name}' not found in the model.")
return None
def explain_image(model_name, image_path, cam_method, feature_module):
model = load_model(model_name)
image = process_image(image_path, model)
target_layer = get_target_layer(model, feature_module)
if target_layer is None:
# Fallback to the last feature module or last convolutional layer
feature_info = get_feature_info(model)
if feature_info:
target_layer = get_target_layer(model, feature_info[-1])
print(f"Using last feature module: {feature_info[-1]}")
else:
# Fallback to finding last convolutional layer
for name, module in reversed(list(model.named_modules())):
if isinstance(module, torch.nn.Conv2d):
target_layer = module
print(f"Fallback: Using last convolutional layer: {name}")
break
if target_layer is None:
raise ValueError("Could not find a suitable target layer.")
cam_image = get_cam_image(model, image, target_layer, cam_method)
return cam_image
def update_feature_modules(model_name):
model = load_model(model_name)
feature_modules = get_feature_info(model)
return gr.Dropdown(choices=feature_modules, value=feature_modules[-1] if feature_modules else None)
with gr.Blocks() as demo:
gr.Markdown("# Explainable AI with timm models")
gr.Markdown("Upload an image, select a model, CAM method, and optionally a specific feature module to visualize the explanation.")
with gr.Row():
with gr.Column():
model_dropdown = gr.Dropdown(choices=MODELS, label="Select Model")
image_input = gr.Image(type="filepath", label="Upload Image")
cam_method_dropdown = gr.Dropdown(choices=list(CAM_METHODS.keys()), label="Select CAM Method")
feature_module_dropdown = gr.Dropdown(label="Select Feature Module (optional)")
explain_button = gr.Button("Explain Image")
with gr.Column():
output_image = gr.Image(type="pil", label="Explained Image")
model_dropdown.change(fn=update_feature_modules, inputs=[model_dropdown], outputs=[feature_module_dropdown])
explain_button.click(
fn=explain_image,
inputs=[model_dropdown, image_input, cam_method_dropdown, feature_module_dropdown],
outputs=[output_image]
)
demo.launch() |