类激活图简介

1import torch
2import torchvision
3import torchvision.transforms as T
4from torchvision.models import mobilenet_v2
5from PIL import Image
6import json
7import matplotlib
8import matplotlib.pyplot as plt
9import numpy as np
10
11class_idx = json.load(open("imagenet_class_index.json", 'r')) # ImageNet ID 到标签的映射
12idx2label = [class_idx[str(k)][1] for k in range(len(class_idx))] 
13
14original_image = Image.open("cat.jpg").resize((224,224)) # 原始图像
15transform = T.Compose([
16    T.Resize((224,224)),
17    T.ToTensor(),
18    T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
19])
20
21image = transform(original_image) # 对图像进行变换和归一化
22image.requires_grad = True 
23torchvision.utils.save_image(image, "./output/image.png", normalize=True)
24
25model = mobilenet_v2() # 加载 MobileNet-V2 模型
26ckpt = torch.load("./mobilenet_v2-b0353104.pth", map_location="cpu")
27model.load_state_dict(ckpt)
28model.eval()
29
30logits = model(image.unsqueeze(0))[0] # 输出 logits，前向传播
31prediction = torch.argmax(logits) # 网络预测结果
32print("模型预测: ", idx2label[prediction.item()])
33
34weight = model.classifier[1].weight.data[prediction.item()]
35
36features = model.features(image.unsqueeze(0))
37features = torch.nn.functional.interpolate(features, (image.shape[1], image.shape[2]), mode='bilinear')
38c, h, w = features.shape[1], features.shape[2], features.shape[3]
39features = features.reshape(c, -1)
40
41weighted_features = torch.matmul(weight, features)
42cam = weighted_features.reshape(h, w)
43
44cam = cam - cam.min()
45cam = cam / cam.max()
46cmap = matplotlib.colormaps.get_cmap('jet')
47cam = cmap(cam.detach().numpy())[:, :, :3]
48heatmap = cam * 0.4 + 0.6 * np.array(original_image) / 255.
49
50plt.imshow(heatmap) # 绘制显著性图
51plt.axis("off")
52plt.tight_layout()
53plt.savefig("./output/saliency.png")