当前位置：首页 > news >正文

pytorch生成CAM热力图-单张图像

news 来源：原创 2025/4/21 14:43:35

利用ImageNet预训练模型生成CAM热力图-单张图像

一、环境搭建
二、主要代码
三、结果展示

代码和图片等资源均来源于哔哩哔哩up主：同济子豪兄
讲解视频：CAM可解释性分析-算法讲解

一、环境搭建

1，安装所需的包

pip install numpy pandas matplotlib requests tqdm opencv-python pillow -i https://pypi.tuna.tsinghua.edu.cn/simple

2，安装 Pytorch

pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113

3，安装 mmcv-full

pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu113/torch1.10.0/index.html

4，下载中文字体文件（用于显示和打印汉字文字）

wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/dataset/SimHei.ttf

5，下载 ImageNet 1000类别信息

wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/dataset/meta_data/imagenet_class_index.csv

6，创建 test_img 文件夹，并下载测试图像到该文件夹

import os
os.mkdir('test_img')

wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/test/border-collie.jpg -P test_img
wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/test/cat_dog.jpg -P test_img
wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/test/0818/room_video.mp4 -P test_img

7，下载安装 torchcam

git clone https://github.com/frgfm/torch-cam.git
pip install -e torch-cam/.

二、主要代码

from PIL import Image

import torch
# 有 GPU 就用 GPU，没有就用 CPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
print('device', device)

# 导入ImageNet预训练模型
from torchvision.models import resnet18
model = resnet18(pretrained=True).eval().to(device)

# 导入自己训练的模型
# model = torch.load('自己训练的模型.pth')
# model = model.eval().to(device)

# 可解释性分析方法有：CAM GradCAM GradCAMpp ISCAM LayerCAM SSCAM ScoreCAM SmoothGradCAMpp XGradCAM

# 方法一：导入可解释性分析方法SmoothGradCAMpp
# from torchcam.methods import SmoothGradCAMpp 
# cam_extractor = SmoothGradCAMpp(model)

# 方法二：导入可解释性分析方法GradCAM
from torchcam.methods import GradCAM
target_layer = model.layer4[-1]    # 选择目标层
cam_extractor = GradCAM(model, target_layer)

# 图片预处理
from torchvision import transforms
# 测试集图像预处理-RCTN：缩放、裁剪、转 Tensor、归一化
test_transform = transforms.Compose([transforms.Resize(256),
                                     transforms.CenterCrop(224),
                                     transforms.ToTensor(),
                                     transforms.Normalize(
                                         mean=[0.485, 0.456, 0.406], 
                                         std=[0.229, 0.224, 0.225])
                                    ])

# 图片分类预测
img_path = 'test_img/border-collie.jpg'
img_pil = Image.open(img_path)
input_tensor = test_transform(img_pil).unsqueeze(0).to(device) # 预处理
pred_logits = model(input_tensor)
# topk()方法用于返回输入数据中特定维度上的前k个最大的元素
pred_top1 = torch.topk(pred_logits, 1)
# pred_id 为图片所属分类对应的索引号，分类和索引号存储在imagenet_class_index.csv
pred_id = pred_top1[1].detach().cpu().numpy().squeeze().item()

# 生成可解释性分析热力图
activation_map = cam_extractor(pred_id, pred_logits)
activation_map = activation_map[0][0].detach().cpu().numpy()

# 可视化
from torchcam.utils import overlay_mask

# overlay_mask 用于构建透明的叠加层
# fromarray 实现array到image的转换
result = overlay_mask(img_pil, Image.fromarray(activation_map), alpha=0.7)

# 为图片添加中文类别显示

# 载入ImageNet 1000 类别中文释义
import pandas as pd
df = pd.read_csv('imagenet_class_index.csv')
idx_to_labels = {}
idx_to_labels_cn = {}
for idx, row in df.iterrows():
    idx_to_labels[row['ID']] = row['class']
    idx_to_labels_cn[row['ID']] = row['Chinese']

# 显示所有中文类别
# idx_to_labels_cn

# 可视化热力图的类别ID，如果为 None，则为置信度最高的预测类别ID
# show_class_id = 231		# 例如 牧羊犬：231 虎猫：281
show_class_id = None

# 可视化热力图的类别ID，如果不指定，则为置信度最高的预测类别ID
if show_class_id:
    show_id = show_class_id
else:
    show_id = pred_id
    show_class_id = pred_id

# 是否显示中文类别
Chinese = True
# Chinese = False

from PIL import ImageDraw
# 在图像上写字
draw = ImageDraw.Draw(result)

if Chinese:
    # 在图像上写中文
    text_pred = 'Pred Class: {}'.format(idx_to_labels_cn[pred_id])
    text_show = 'Show Class: {}'.format(idx_to_labels_cn[show_class_id])
else:
    # 在图像上写英文
    text_pred = 'Pred Class: {}'.format(idx_to_labels[pred_id])
    text_show = 'Show Class: {}'.format(idx_to_labels[show_class_id])

from PIL import ImageFont, ImageDraw
# 导入中文字体，指定字体大小
font = ImageFont.truetype('SimHei.ttf', 30)

# 文字坐标，中文字符串，字体，rgba颜色
draw.text((10, 10), text_pred, font=font, fill=(255, 0, 0, 1))
draw.text((10, 50), text_show, font=font, fill=(255, 0, 0, 1))

#输出结果图
result

注意：

可解释性方法的选择有多种，代码中提供了 SmoothGradCAMpp 和 GradCAM 两种方法；
模型选择也有pytorch预训练模型和自己训练的模型两种，代码中演示了 ImageNet图像分类模型，图片类别文件为 imagenet_class_index.csv；若为自己的模型则还需要修改 “为图片载入类别的部分代码”