CAM（热力图）Torch代码实现

在做图像处理时，经常会根据热力图进行分析（如下图所示）：

---

                  原图                                           热力图

代码：

import io
import os
import requests
from PIL import Image
from torchvision import models, transforms
from torch.autograd import Variable
from torch.nn import functional as F
import numpy as np
import cv2
import pdb
from IC_DenseNet161 import densenet161

# input image
image_path = './data/test/0050/5755.jpg'

# networks such as googlenet, resnet, densenet already use global average pooling at the end, so CAM could be used directly.

net = densenet161()
finalconv_name = 'features'

net.eval()

# hook the feature extractor
features_blobs = []
def hook_feature(module, input, output):
    features_blobs.append(output.data.cpu().numpy())

net._modules.get(finalconv_name).register_forward_hook(hook_feature)

# get the softmax weight
# 倒数第二层
params = list(net.parameters())
weight_softmax = np.squeeze(params[-2].data.numpy())

def returnCAM(feature_conv, weight_softmax, class_idx):
    # generate the class activation maps upsample to 256x256
    size_upsample = (224, 224)
    bz, nc, h, w = feature_conv.shape
    output_cam = []
    for idx in class_idx:
        # 回到GAP的值
        cam = weight_softmax[idx].dot(feature_conv.reshape((nc, h*w)))
        cam = cam.reshape(h, w)
        # np.min 返回数组的最小值或沿轴的最小值。
        cam = cam - np.min(cam)
        cam_img = cam / np.max(cam)
        #  np.uint8(） Create a data type object.
        cam_img = np.uint8(255 * cam_img)
        output_cam.append(cv2.resize(cam_img, size_upsample))
    return output_cam


normalize = transforms.Normalize(
   mean=[0.485, 0.456, 0.406],
   std=[0.229, 0.224, 0.225]
)
preprocess = transforms.Compose([
   # transforms.Resize((224,224)),
   transforms.ToTensor(),
   # normalize
])

image_path = os.path.expanduser(image_path)
img_pil = Image.open(image_path)
img_pil.save('test.jpg')

img_tensor = preprocess(img_pil)
img_variable = Variable(img_tensor.unsqueeze(0))
logit = net(img_variable)

h_x = F.softmax(logit, dim=1).data.squeeze()
probs, idx = h_x.sort(0, True)
probs = probs.numpy()
idx = idx.numpy()

# generate class activation mapping for the top1 prediction
CAMs = returnCAM(features_blobs[0], weight_softmax, [idx[0]])

# render the CAM and output
# print('output CAM.jpg for the top1 prediction: %s'%classes[idx[0]])
print('output CAM.jpg ')
img = cv2.imread('./CAM/test.jpg')
height, width, _ = img.shape
heatmap = cv2.applyColorMap(cv2.resize(CAMs[0],(width, height)), cv2.COLORMAP_JET)
result = heatmap * 0.3 + img * 0.5
cv2.imwrite('./CAM/CAM.jpg', result)