python
复制代码
import torch
import torch.nn.functional as F
import numpy as np
import cv2
import matplotlib.pyplot as plt
from torchvision import models, transforms
from PIL import Image
# 检查是否有 GPU
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
class GradCAM:
def __init__(self, model, target_layer):
"""
初始化 Grad-CAM
:param model: 已加载权重的模型
:param target_layer: 需要可视化的目标卷积层 (通常是最后一个卷积层)
"""
self.model = model.eval().to(device) # 设为评估模式
self.target_layer = target_layer
self.gradients = None
self.activations = None
# 注册 Hook
self._register_hooks()
def _register_hooks(self):
"""
注册前向和反向传播的 Hook
"""
# Hook 1: 前向传播时，保存特征图 (Activations)
def save_activation(module, input, output):
self.activations = output.detach()
# Hook 2: 反向传播时，保存梯度 (Gradients)
def save_gradient(module, grad_input, grad_output):
# grad_output 是一个 tuple，通常取第一项
self.gradients = grad_output[0].detach()
# 将 Hook 注册到目标层
self.target_layer.register_forward_hook(save_activation)
self.target_layer.register_full_backward_hook(save_gradient)
def generate(self, input_image, target_class_index=None):
"""
生成 Grad-CAM 热力图
:param input_image: 预处理后的输入图像张量 (1, C, H, W)
:param target_class_index: 目标类别的索引 (如果为 None，则自动选择预测概率最高的类)
"""
# 1. 前向传播
output = self.model(input_image)
# 如果未指定类别，取预测概率最大的类别
if target_class_index is None:
target_class_index = torch.argmax(output, dim=1).item()
# 2. 构造目标信号 (One-hot 形式，仅对目标类别求导)
self.model.zero_grad()
one_hot_output = torch.zeros_like(output)
one_hot_output[0][target_class_index] = 1
# 3. 反向传播 (触发 Backward Hook，获取梯度)
output.backward(gradient=one_hot_output, retain_graph=True)
# --- 以下是 Grad-CAM 的数学计算过程 ---
# 获取捕获的梯度和特征图
grads = self.gradients # [Batch, Channel, Height, Width]
fmap = self.activations # [Batch, Channel, Height, Width]
# Step A: 全局平均池化 (GAP) 计算权重 alpha
# 对每个通道的梯度求均值
weights = torch.mean(grads, dim=(2, 3), keepdim=True)
# Step B: 加权组合特征图
# weights * fmap 利用广播机制
cam = torch.sum(weights * fmap, dim=1, keepdim=True)
# Step C: ReLU 激活 (只保留对类别有正向贡献的特征)
cam = F.relu(cam)
# Step D: 归一化并调整大小
cam = cam.squeeze().cpu().numpy() # 转为 numpy
if np.max(cam) == 0: # 防止全0导致除0错误
return np.zeros((input_image.shape[2], input_image.shape[3]))
cam = (cam - np.min(cam)) / (np.max(cam) - np.min(cam)) # 归一化到 0-1
# 将热力图缩放到原始图像大小
original_h, original_w = input_image.shape[2], input_image.shape[3]
cam = cv2.resize(cam, (original_w, original_h))
return cam
def show_cam_on_image(self, img_path, cam, alpha=0.5):
"""
可视化：将热力图叠加到原图上
"""
img = cv2.imread(img_path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = np.float32(img) / 255
# 调整 cam 大小确保匹配 (cv2.resize 可能会有微小误差，强制 resize)
heatmap = cv2.resize(cam, (img.shape[1], img.shape[0]))
# 应用伪彩色 (JET colormap)
heatmap = np.uint8(255 * heatmap)
heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
heatmap = cv2.cvtColor(heatmap, cv2.COLOR_BGR2RGB)
heatmap = np.float32(heatmap) / 255
# 叠加
cam_result = heatmap * alpha + img * (1 - alpha)
cam_result = cam_result / np.max(cam_result)
plt.figure(figsize=(10, 5))
plt.subplot(1, 2, 1)
plt.imshow(img)
plt.title("Original Image")
plt.subplot(1, 2, 2)
plt.imshow(cam_result)
plt.title("Grad-CAM")
plt.show()