使用GradCam得到测试图像的热力图

  • 配置环境,安装必要的包!


        1.先激活你的conda环境!         

        2.安装 grad-cam包:

pip install grad-cam

        3.安装 torchcam包:

pip install torchcam

        4.下载 pytorch-grad-cam:

git clone https://github.com/jacobgil/pytorch-grad-cam.git

 处理流程

  1.  指定设备;
  2. 加载模型断点;
  3.  加载模型中最后一个卷积层的权重参数;
  4. 加载并预处理测试样本;
  5. 初始化并使用Grad-Cam;
  6. 保存热力图;

 代码

import os,sys
#解决linux下导报无法识别的问题
BASE_DIR = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, BASE_DIR)
from pytorch_grad_cam import GradCAM
import torch
# 注意:在此处导入你自己的模型
from model_visual import FaceModel
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
from torchcam.utils import overlay_mask
import numpy as np
from PIL import Image
from torchvision import transforms

def get_last_conv(m):
    """
    Get the last conv layer in an Module.
    """
    convs = filter(lambda k: isinstance(k, torch.nn.Conv2d), m.modules())
    return list(convs)[-1]

#################### 指定设备 ####################
# 指定GPU 或 CPU
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')

#################### 加载模型断点 ####################
# 初始化模型
model = FaceModel(isTrain = False).to(device)
# 指定保存的模型断点路径
model_path = "xxx/best/best_xxx.pth"
# 加载断点
checkpoint = torch.load(model_path)
# 断点加载到模型上
model.load_state_dict(checkpoint)
# 开始进入评估模式
model.eval()

#################### 加载最后一个卷积层的权重参数 ####################
layers = get_last_conv(model)
target_layers = [layers]

#################### 加载并预处理测试样本 ####################
# 用于预处理测试样本
test_transform = transforms.Compose([
        transforms.Resize((256, 256), Image.BICUBIC),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.59416118, 0.51189164, 0.45280306],
                            std=[0.25687563, 0.26251543, 0.26231294])
        ])

# 指定待测图像路径
img_path = "images/a_junkai/OULU/f_Test_256/Test_files_1_1/1_1_54_1/1.png"
# 根据RGB路径,来读取RGB图像
rgb_img = Image.open(img_path).convert("RGB")
rgb_tensor = test_transform(rgb_img).unsqueeze(0).to(device)
# 预测category_id
idx = np.argmax(model.cls_feat.detach().cpu().numpy())  # predict id
target_label = [ClassifierOutputTarget(idx)]

#################### 初始化并使用Grad-Cam ####################
# 初始化GradCAM模型
cam = GradCAM(model=model, target_layers=target_layers, use_cuda=torch.cuda.is_available())
cam_map = cam(input_tensor= rgb_tensor, targets= target_label)[0]  # 不加平滑
result_rgb = overlay_mask(rgb_img, Image.fromarray(cam_map), alpha=0.6)  # alpha越小,原图越淡

#################### 指定保存路径并保存热力图 ####################
save_depthPath = "p1/heatmap/rgb/"

if not os.path.exists(save_depthPath):
    os.makedirs(save_depthPath)

save_rgbImg = os.path.join(save_depthPath,"test_visual.jpg")
# 保存热力图
result_rgb.save(save_rgbImg)

 Visualization

featmap

参考:

https://github.com/TommyZihao/Train_Custom_Dataset/tree/main/

猜你喜欢

转载自blog.csdn.net/power_kaikaige/article/details/129468502