利用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 Imageimport 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_idshow_class_id = pred_id# 是否显示中文类别
Chinese = True
# Chinese = Falsefrom 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;若为自己的模型则还需要修改 “为图片载入类别的部分代码”