2021-11-02

标题类激活图（CAM）实现

参考文章：类激活图（CAM）代码+原理详解【pytorch亲测有效】
在此基础上去除了预测信息。解释都在代码中：

from __future__ import print_function, division
import os
from PIL import Image
import matplotlib.pyplot as plt
import torch
from model_SERes_2 import resnet50
from torchvision import models, transforms
import cv2
import numpy as np
from tqdm import tqdm


device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
# 数据处理部分
data_transform = transforms.Compose(
        [transforms.Resize(256),
         transforms.CenterCrop(224),
         transforms.ToTensor(),
         transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])

# 加载模型和权重
model = resnet50(num_classes=30)
weights_path = "./resNet50.pth"
assert os.path.exists(weights_path), "file: '{}' dose not exist.".format(weights_path)
model.load_state_dict(torch.load(weights_path, map_location=device))

# 这部分不怎么懂
def returnCAM(feature_conv, weight_softmax, class_idx):
    bz, nc, h, w = feature_conv.shape  # 1,2048,7,7
    output_cam = []
    for idx in class_idx:  # 若只输出预测概率最大值结果不需要for循环
        feature_conv = feature_conv.reshape((nc, h * w))
        cam = weight_softmax[idx].dot(
            feature_conv.reshape((nc, h * w)).detach().numpy()) # (2048, ) * (2048, 7*7) -> (7*7, ) （n,）是一个数组，既不是行向量也不是列向量
        cam = cam.reshape(h, w)
        cam_img = (cam - cam.min()) / (cam.max() - cam.min())  # Normalize
        cam_img = np.uint8(255 * cam_img)  # Format as CV_8UC1 (as applyColorMap required)
        output_cam.append(cam_img)
    return output_cam


# get weight matrix of full connection
fc_weights = model.state_dict()['fc.weight'].cpu().numpy()  # [2,2048]


model.eval()

"""
对一个文件夹内的图片进行处理
path:当前路径下存放待处理图片的文件夹，图片从1顺序命名
save_img：文件保存文件夹
"""
path = "05"
imgs_list = os.listdir(path)
imgs_list.sort(key=lambda x: int(x[:-4]))
imgs_path = [os.path.join(path, x) for x in imgs_list]
val_bar = tqdm(imgs_path)
save_img = "./save4_picture"
if not os.path.exists(save_img):
    os.mkdir(save_img)

for index, img_path in enumerate(val_bar, start=1):
    # 图片保存路径
    save_path = save_img + "/cam" + str(index) + ".jpg"

    img = Image.open(img_path).convert('RGB')
    img_tensor = data_transform(img).unsqueeze(0)  # [1,3,224,224]
    inputs = img_tensor

    logit = model(inputs)  # [1,2] -> [ 3.3207, -2.9495]
    h_x = torch.nn.functional.softmax(logit, dim=1).data.squeeze()  # tensor([0.9981, 0.0019])
    probs, idx = h_x.sort(0, True)  # sorted in descending order

    probs = probs.cpu().numpy()  # if tensor([0.0019,0.9981]) ->[0.9981, 0.0019]
    idx = idx.cpu().numpy()  # [1, 0]

    # 得到最后一层卷积层
    img_1 = model.conv1(img_tensor)
    img_1 = model.bn1(img_1)
    img_1 = model.relu(img_1)
    img_1 = model.maxpool(img_1)
    img_1 = model.layer1(img_1)
    img_1 = model.layer2(img_1)
    img_1 = model.layer3(img_1)
    features = model.layer4(img_1)

    # features = model.finalconv.cpu().numpy()  # [1,2048,7,7]

    CAMs = returnCAM(features, fc_weights, [idx[0]])  # output the most probability class activate map

    img = cv2.imread(img_path)
    height, width, _ = img.shape  # get input image size
    heatmap = cv2.applyColorMap(cv2.resize(CAMs[0], (width, height)),
                                cv2.COLORMAP_JET)  # CAM resize match input image size
    heatmap[np.where(CAMs[0] <= 100)] = 0
    result = heatmap * 0.3 + img * 0.5  # ratio

    cv2.imwrite(save_path, result)

    val_bar.desc = "done"
    # print("*****done:{index}*****".format(index=index), end="\r")
print()
print("Finishing")

说明

我模型前向传播部分如下，我模型训练好了，在模型中加入self.finalconv = x.detach()，训练好的权重不能用，所以我用如下得到最后一层卷积层。

img_1 = model.conv1(img_tensor)
    img_1 = model.bn1(img_1)
    img_1 = model.relu(img_1)
    img_1 = model.maxpool(img_1)
    img_1 = model.layer1(img_1)
    img_1 = model.layer2(img_1)
    img_1 = model.layer3(img_1)
    features = model.layer4(img_1)