torchcam-single-ImageNet4

# 导入工具包
import matplotlib.pyplot as plt
from PIL import Image
import torch
from torchvision.models import resnet18
from torchcam.methods import SmoothGradCAMpp
# CAM GradCAM GradCAMpp ISCAM LayerCAM SSCAM ScoreCAM SmoothGradCAMpp XGradCAM
from torchvision import transforms
from torchcam.utils import overlay_mask
import pandas as pd
from PIL import ImageFont, ImageDraw

def torchcam_single_ImageNet4(image_path,csv_file,show_class_id,Chinese,font):
    # 有 GPU 就用 GPU，没有就用 CPU
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print('device', device)
    # 导入ImageNet预训练模型
    model = resnet18(pretrained=True).eval().to(device)
    # 导入可解释性分析方法
    cam_extractor = SmoothGradCAMpp(model)
    # 测试集图像预处理-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])])

    # 整理代码：设置类别、中文类别显示
    # 载入ImageNet 1000图像分类标签
    # ImageNet 1000类别中文释义：https://github.com/ningbonb/imagenet_classes_chinese
    df = pd.read_csv(csv_file)
    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']

    show_class_id=show_class_id
    Chinese=Chinese
    font=font

    img_pil = Image.open(image_path)
    input_tensor = test_transform(img_pil).unsqueeze(0).to(device)  # 预处理
    # 对输入张量进行模型推断
    pred_logits = model(input_tensor)  # 使用模型进行推断，获取预测 logits（未经 softmax 处理的输出）
    # 获取最可能的类别标签
    pred_top1 = torch.topk(pred_logits, 1)  # 从预测 logits 中获取最高的概率及对应的索引
    pred_id = pred_top1[1].detach().cpu().numpy().squeeze().item()  # 将预测结果从张量转换为 NumPy 数组，并提取最可能的类别索引
    # 可视化热力图的类别ID，如果不指定，则为置信度最高的预测类别ID
    if show_class_id:
        show_id = show_class_id
    else:
        show_id = pred_id
    # 生成可解释性分析热力图
    activation_map = cam_extractor(show_id, pred_logits)  # 使用 cam_extractor 函数生成类激活图，传入预测的类别索引和预测 logits
    # 处理激活图结果
    activation_map = activation_map[0][0].detach().cpu().numpy()  # 从张量中提取类激活图的 NumPy 数组表示
    result = overlay_mask(img_pil, Image.fromarray(activation_map), alpha=0.5)# 用于将类激活图（Activation Map）叠加到原始图像上，以可视化模型对输入图像的关注区域。
    # 在图像上写字
    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])
    # 文字坐标，中文字符串，字体，rgba颜色
    draw.text((50, 100), text_pred, font=font, fill=(255, 0, 0, 1))
    draw.text((50, 200), text_show, font=font, fill=(255, 0, 0, 1))
    # 可视化
    plt.subplots(figsize=(6, 6))
    plt.imshow(result)
    plt.show()


if __name__ == '__main__':
    image_path='test_img/cat_dog.jpg'
    csv_file_path='imagenet_class_index.csv'
    # 可视化热力图的类别ID，如果为 None，则为置信度最高的预测类别ID
    show_class_id = 231
    # 是否显示中文类别
    Chinese = True # False
    # 导入中文字体，指定字体大小
    font_path='SimHei.ttf'
    font_size=80
    font = ImageFont.truetype(font_path, font_size)
    torchcam_single_ImageNet4(image_path,
                              csv_file_path,
                              show_class_id,
                              Chinese,
                              font)