• pytorch生成CAM热力图-单张图像


    利用ImageNet预训练模型生成CAM热力图-单张图像

    代码和图片等资源均来源于哔哩哔哩up主:同济子豪兄
    讲解视频:CAM可解释性分析-算法讲解

    一、环境搭建

    1,安装所需的包

    pip install numpy pandas matplotlib requests tqdm opencv-python pillow -i https://pypi.tuna.tsinghua.edu.cn/simple
    
    • 1

    2,安装 Pytorch

    pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113
    
    • 1

    3,安装 mmcv-full

    pip install mmcv-full -f https://download.openmmlab.com/mmcv/dist/cu113/torch1.10.0/index.html
    
    • 1

    4,下载中文字体文件(用于显示和打印汉字文字)

    wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/dataset/SimHei.ttf
    
    • 1

    5,下载 ImageNet 1000类别信息

    wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/dataset/meta_data/imagenet_class_index.csv
    
    • 1

    6,创建 test_img 文件夹,并下载测试图像到该文件夹

    import os
    os.mkdir('test_img')
    
    wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/test/border-collie.jpg -P test_img
    wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/test/cat_dog.jpg -P test_img
    wget https://zihao-openmmlab.obs.cn-east-3.myhuaweicloud.com/20220716-mmclassification/test/0818/room_video.mp4 -P test_img
    
    • 1
    • 2
    • 3
    • 4
    • 5
    • 6

    7,下载安装 torchcam

    git clone https://github.com/frgfm/torch-cam.git
    pip install -e torch-cam/.
    
    • 1
    • 2

    二、主要代码

    from PIL import Image
    
    import torch
    # 有 GPU 就用 GPU,没有就用 CPU
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    print('device', device)
    
    # 导入ImageNet预训练模型
    from torchvision.models import resnet18
    model = resnet18(pretrained=True).eval().to(device)
    
    # 导入自己训练的模型
    # model = torch.load('自己训练的模型.pth')
    # model = model.eval().to(device)
    
    # 可解释性分析方法有:CAM GradCAM GradCAMpp ISCAM LayerCAM SSCAM ScoreCAM SmoothGradCAMpp XGradCAM
    
    # 方法一:导入可解释性分析方法SmoothGradCAMpp
    # from torchcam.methods import SmoothGradCAMpp 
    # cam_extractor = SmoothGradCAMpp(model)
    
    # 方法二:导入可解释性分析方法GradCAM
    from torchcam.methods import GradCAM
    target_layer = model.layer4[-1]    # 选择目标层
    cam_extractor = GradCAM(model, target_layer)
    
    # 图片预处理
    from torchvision import transforms
    # 测试集图像预处理-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])
                                        ])
    
    # 图片分类预测
    img_path = 'test_img/border-collie.jpg'
    img_pil = Image.open(img_path)
    input_tensor = test_transform(img_pil).unsqueeze(0).to(device) # 预处理
    pred_logits = model(input_tensor)
    # topk()方法用于返回输入数据中特定维度上的前k个最大的元素
    pred_top1 = torch.topk(pred_logits, 1)
    # pred_id 为图片所属分类对应的索引号,分类和索引号存储在imagenet_class_index.csv
    pred_id = pred_top1[1].detach().cpu().numpy().squeeze().item()
    
    # 生成可解释性分析热力图
    activation_map = cam_extractor(pred_id, pred_logits)
    activation_map = activation_map[0][0].detach().cpu().numpy()
    
    # 可视化
    from torchcam.utils import overlay_mask
    
    # overlay_mask 用于构建透明的叠加层
    # fromarray 实现array到image的转换
    result = overlay_mask(img_pil, Image.fromarray(activation_map), alpha=0.7)
    
    # 为图片添加中文类别显示
    
    # 载入ImageNet 1000 类别中文释义
    import pandas as pd
    df = pd.read_csv('imagenet_class_index.csv')
    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']
    
    # 显示所有中文类别
    # idx_to_labels_cn
    
    # 可视化热力图的类别ID,如果为 None,则为置信度最高的预测类别ID
    # show_class_id = 231		# 例如 牧羊犬:231 虎猫:281
    show_class_id = None
    
    # 可视化热力图的类别ID,如果不指定,则为置信度最高的预测类别ID
    if show_class_id:
        show_id = show_class_id
    else:
        show_id = pred_id
        show_class_id = pred_id
    
    # 是否显示中文类别
    Chinese = True
    # Chinese = False
    
    from PIL import ImageDraw
    # 在图像上写字
    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])
    
    from PIL import ImageFont, ImageDraw
    # 导入中文字体,指定字体大小
    font = ImageFont.truetype('SimHei.ttf', 30)
    
    # 文字坐标,中文字符串,字体,rgba颜色
    draw.text((10, 10), text_pred, font=font, fill=(255, 0, 0, 1))
    draw.text((10, 50), text_show, font=font, fill=(255, 0, 0, 1))
    
    #输出结果图
    result
    
    • 1
    • 2
    • 3
    • 4
    • 5
    • 6
    • 7
    • 8
    • 9
    • 10
    • 11
    • 12
    • 13
    • 14
    • 15
    • 16
    • 17
    • 18
    • 19
    • 20
    • 21
    • 22
    • 23
    • 24
    • 25
    • 26
    • 27
    • 28
    • 29
    • 30
    • 31
    • 32
    • 33
    • 34
    • 35
    • 36
    • 37
    • 38
    • 39
    • 40
    • 41
    • 42
    • 43
    • 44
    • 45
    • 46
    • 47
    • 48
    • 49
    • 50
    • 51
    • 52
    • 53
    • 54
    • 55
    • 56
    • 57
    • 58
    • 59
    • 60
    • 61
    • 62
    • 63
    • 64
    • 65
    • 66
    • 67
    • 68
    • 69
    • 70
    • 71
    • 72
    • 73
    • 74
    • 75
    • 76
    • 77
    • 78
    • 79
    • 80
    • 81
    • 82
    • 83
    • 84
    • 85
    • 86
    • 87
    • 88
    • 89
    • 90
    • 91
    • 92
    • 93
    • 94
    • 95
    • 96
    • 97
    • 98
    • 99
    • 100
    • 101
    • 102
    • 103
    • 104
    • 105
    • 106
    • 107
    • 108
    • 109
    • 110

    注意:

    1. 可解释性方法的选择有多种,代码中提供了 SmoothGradCAMpp 和 GradCAM 两种方法;
    2. 模型选择也有pytorch预训练模型和自己训练的模型两种,代码中演示了 ImageNet图像分类 模型,图片类别文件为 imagenet_class_index.csv;若为自己的模型则还需要修改 “为图片载入类别的部分代码”

    三、结果展示

  • 相关阅读:
    Python Selenium元素定位方法详解
    【课程】SP Module4 音频滤波器
    系统架构设计高级技能 · 构件与中间件技术
    Redis的安装
    解决这三大问题,运维效率将超90%的医院
    Spring Session原理解析
    生物化学 核磁共振 氢谱 n+1定律 邻碳耦合 同碳耦合
    软考 系统架构设计师 简明教程 | 系统运行与软件维护
    三子棋——C语言初阶
    Golang Testify介绍
  • 原文地址:https://blog.csdn.net/weixin_44624410/article/details/132899013