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

利用ImageNet预训练模型生成CAM热力图-单张图像 一、环境搭建二、主要代码三、结果展示

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

一、环境搭建

1，安装所需的包

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

2，安装 Pytorch

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

3，安装 mmcv-full

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

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

wget zihao-openmmlab.obs -east-3.myhuaweicloud /20220716-mmclassification/dataset/SimHei.ttf

5，下载 ImageNet 1000类别信息

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

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

import os os.mkdir('test_img') wget zihao-openmmlab.obs -east-3.myhuaweicloud /20220716-mmclassification/test/border-collie.jpg -P test_img wget zihao-openmmlab.obs -east-3.myhuaweicloud /20220716-mmclassification/test/cat_dog.jpg -P test_img wget zihao-openmmlab.obs -east-3.myhuaweicloud /20220716-mmclassification/test/0818/room_video.mp4 -P test_img

7，下载安装 torchcam

git clone github /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；若为自己的模型则还需要修改 “为图片载入类别的部分代码” 三、结果展示