day42打卡

@浙大疏锦行

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