Spaces:
Running
Running
File size: 3,885 Bytes
af59988 |
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 137 138 |
"""
Grad-CAM visualization for model interpretability.
"""
import torch
import numpy as np
from PIL import Image
from pathlib import Path
from typing import Union
import matplotlib.pyplot as plt
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import show_cam_on_image
from .dataset import get_transforms
from .config import IMAGENET_MEAN, IMAGENET_STD, CLASS_NAMES
def get_gradcam(model, target_layer=None):
"""Create GradCAM object for the model."""
if target_layer is None:
# Use the last conv layer of EfficientNet
target_layer = model.backbone.features[-1]
return GradCAM(model=model, target_layers=[target_layer])
def denormalize_image(tensor: torch.Tensor) -> np.ndarray:
"""Denormalize tensor to numpy image [0,1]."""
mean = torch.tensor(IMAGENET_MEAN).view(3, 1, 1)
std = torch.tensor(IMAGENET_STD).view(3, 1, 1)
img = tensor.cpu() * std + mean
img = img.permute(1, 2, 0).numpy()
return np.clip(img, 0, 1)
def generate_gradcam(
model,
image: Union[str, Path, Image.Image],
device: torch.device
) -> tuple:
"""Generate Grad-CAM heatmap for an image."""
model.eval()
# Load and transform image
if isinstance(image, (str, Path)):
image = Image.open(image).convert('RGB')
transform = get_transforms(is_training=False)
img_tensor = transform(image).unsqueeze(0).to(device)
# Get prediction
with torch.no_grad():
output = model(img_tensor)
prob = torch.sigmoid(output).item()
pred_class = CLASS_NAMES[1] if prob > 0.5 else CLASS_NAMES[0]
confidence = prob if prob > 0.5 else 1 - prob
# Generate Grad-CAM
cam = get_gradcam(model)
grayscale_cam = cam(input_tensor=img_tensor, targets=None)[0]
# Create visualization
rgb_img = denormalize_image(img_tensor[0])
cam_image = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)
return cam_image, pred_class, confidence, rgb_img
def plot_gradcam(
model,
image_path: Union[str, Path],
true_label: str,
device: torch.device,
save_path: str = None
):
"""Plot original image with Grad-CAM overlay."""
cam_image, pred_class, confidence, original = generate_gradcam(model, image_path, device)
fig, axes = plt.subplots(1, 2, figsize=(10, 4))
# Original
axes[0].imshow(original)
axes[0].set_title(f"Original\nTrue: {true_label}")
axes[0].axis('off')
# Grad-CAM
color = 'green' if pred_class == true_label else 'red'
axes[1].imshow(cam_image)
axes[1].set_title(f"Grad-CAM\nPred: {pred_class} ({confidence:.1%})", color=color)
axes[1].axis('off')
plt.tight_layout()
if save_path:
plt.savefig(save_path, dpi=150, bbox_inches='tight')
plt.show()
return pred_class, confidence
def plot_gradcam_grid(
model,
image_paths: list,
true_labels: list,
device: torch.device,
save_path: str = None,
title: str = "Grad-CAM Visualizations"
):
"""Plot grid of Grad-CAM visualizations."""
n = len(image_paths)
fig, axes = plt.subplots(n, 2, figsize=(8, 3 * n))
if n == 1:
axes = axes.reshape(1, -1)
for i, (path, true_label) in enumerate(zip(image_paths, true_labels)):
cam_image, pred_class, confidence, original = generate_gradcam(model, path, device)
# Original
axes[i, 0].imshow(original)
axes[i, 0].set_title(f"True: {true_label}")
axes[i, 0].axis('off')
# Grad-CAM
color = 'green' if pred_class == true_label else 'red'
axes[i, 1].imshow(cam_image)
axes[i, 1].set_title(f"Pred: {pred_class} ({confidence:.1%})", color=color)
axes[i, 1].axis('off')
plt.suptitle(title, fontsize=14, fontweight='bold')
plt.tight_layout()
if save_path:
plt.savefig(save_path, dpi=150, bbox_inches='tight')
plt.show()
|