-
配置环境,安装必要的包!
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
处理流程
- 指定设备;
- 加载模型断点;
- 加载模型中最后一个卷积层的权重参数;
- 加载并预处理测试样本;
- 初始化并使用Grad-Cam;
- 保存热力图;
代码
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
参考:
https://github.com/TommyZihao/Train_Custom_Dataset/tree/main/