pyTorchCAM集合（GradCAM,ScoreCAM,SSCAM,ISCAM……）代码实现

基于ResNet18和Kaggle猫狗数据集实现

from torchvision.io.image import read_image
from torchvision.transforms.functional import normalize, resize, to_pil_image
from torchvision.models import resnet18,alexnet
from torchcam.methods import GradCAM, GradCAMpp, ScoreCAM
from torchcam.utils import overlay_mask
import numpy as np
from torchvision import datasets, transforms
import os
from torch.utils.data import Dataset
import glob
from torch.utils.data import DataLoader
import pandas as pd
import cv2
import torch
import torch.nn as nn
from PIL import Image

mean=[0.485, 0.456, 0.406]
std=[0.229, 0.224, 0.225]

class test_ImageDataset(Dataset):
    def __init__(self, root):

        # Transforms for low resolution images and high resolution images
        self.transform = transforms.Compose(
            [
                transforms.Resize((256, 256)),
                transforms.ToTensor(),
                transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
            ])# 归一化

        self.files = sorted(glob.glob(root + "/*.*"))

    def __getitem__(self, index):
        path = self.files[index % len(self.files)]
        label = path[path.index(".")-3:path.index(".")]
        img = Image.open(self.files[index % len(self.files)])
        img = self.transform(img)
        #print(img.shape)
        #print(label)
        if label == "cat":
            label = [0, 1]
        if label == "dog":
            label = [1, 0]
        label = torch.Tensor(label)
        #print(label)
        return img, label
    def __len__(self):
        return len(self.files)

def denormalize(tensors):
    """ Denormalizes image tensors using mean and std """
    for c in range(3):
        tensors[c,:,:].mul_(std[c]).add_(mean[c])
    return torch.clamp(tensors, 0, 255)


# 测试
def test(model, device, test_loader):
    model.eval()
    cam_extractor = ScoreCAM(model, target_layer="encoder") #根据torchcam库修改即可
    #cam_extractor = GradCAM(model, target_layer="encoder")
    index = 0
    for data, target in test_loader:
        data_de = denormalize(data.squeeze())  # 防止后面data变成cuda
        data, target = data.to(device), target.to(device)
        out = model(data)
        pred = out.max(1, keepdim=True)[1].squeeze(1).cpu().numpy()  # 找到概率最大的下标
        # print(target)
        target = target.max(1, keepdim=True)[1].detach().squeeze(1).cpu().numpy()
        bo = (pred == target)
        dir = "./ScoreCAM_output/test/%d_attentionMap_%d_%s" % (index, int(pred[0]), str(bo))
        os.makedirs(dir, exist_ok=True)

        ori__img = np.array(data_de)[::-1, :, :].transpose(1, 2, 0) * 255.0
        ori_img = cv2.resize(ori__img, (500, 500))
        original_img = ori_img.astype(np.uint8)
        cv2.imwrite("%s/%d_ORIImg.jpg" % (dir, index), original_img)

        activation_map = cam_extractor(out.squeeze(0).argmax().item(), out)
        activation_map = activation_map[0].detach().squeeze().cpu().numpy() * 255.0
        activation_map = cv2.resize(activation_map, (500, 500))
        activation_map = activation_map.astype(np.uint8)
        map = cv2.applyColorMap(activation_map, cv2.COLORMAP_JET)
        add_map = cv2.addWeighted(map, 0.7, original_img, 0.3, 0)
        cv2.imwrite("%s/%d_AttImg.jpg" % (dir, index), map)
        cv2.imwrite("%s/%d_AddImg.jpg" % (dir, index), add_map)
        index += 1

if __name__ == '__main__':
    BATCH_SIZE = 1  # 大概需要2G的显存
    DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # 让torch判断是否使用GPU，建议使用GPU环境，因为会快很多

    # 下载测试集
    test_loader = DataLoader(
        test_ImageDataset("./test"),
        batch_size=BATCH_SIZE,
        shuffle=True,
        num_workers=0,
    )
    model = resnet18().to(DEVICE)
    model.fc = nn.Linear(in_features=512,out_features=2,bias=True).to(DEVICE)
    print(model)
    model.load_state_dict(torch.load("./saved_model/AlexNet_PreTrain_ALL_CAT&DOG_iteration_50_lr_0.0001.pth"))

    test(model,DEVICE,test_loader)

torchcam库可以直接用pip install torchcam 进行安装