不好之处就是这种方法不能最大程度还原最后的热图，会产生很多噪声，不过大体区域还是接近的，代码如下：
python
import cv2
import os
import numpy as np
from PIL import Image
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
import sys
import time
import shutil
import argparse
import json
from torch.utils.data import DataLoader
from utils import set_logger, update_lr, get_pck_with_sigma, get_pred_coordinates, save_images, save_limb_images
# ================== Grad-CAM ==================
class GradCAM:
def __init__(self, model, target_layer):
self.model = model
self.target_layer = target_layer
self.gradients = None # 存储梯度
self.activations = None # 存储激活值
# 绑定梯度钩子
self.hook_layers()
def hook_layers(self):
""" 绑定 hook，提取梯度和特征图 """
def forward_hook(module, input, output):
self.activations = output # 存储前向传播的特征图
def backward_hook(module, grad_in, grad_out):
self.gradients = grad_out[0] # 存储反向传播的梯度
self.target_layer.register_forward_hook(forward_hook)
self.target_layer.register_backward_hook(backward_hook)
def forward(self, image, joint_index):
""" 计算 Grad-CAM（针对某个关键点）"""
self.model.eval() # 进入评估模式
image = image.requires_grad_(True) # 需要梯度信息
output = self.model(image) # 形状 (1, num_joints, H, W)
# 选择某个关键点的热图进行梯度计算
heatmap = output[:, joint_index, :, :] # (1, H, W)
loss = heatmap.sum() # 让 loss 与该热图相关
self.model.zero_grad()
loss.backward() # 计算梯度
# 计算 Grad-CAM 权重（对梯度做全局平均池化）
gradients = self.gradients.mean(dim=(2, 3), keepdim=True) # (1, C, 1, 1)
activations = self.activations # (1, C, H, W)
cam = (activations * gradients).sum(dim=1, keepdim=True) # (1, 1, H, W)
cam = F.relu(cam) # 保证非负
cam = cam.squeeze().cpu().detach().numpy() # (H, W)
return cam
def overlay_heatmap(self, image, cam):
""" 叠加 Grad-CAM 热图，使用默认颜色 """
# image 为 BGR 格式
cam = cv2.resize(cam, (image.shape[1], image.shape[0])) # 调整大小
cam = (cam - cam.min()) / (cam.max() - cam.min() + 1e-8) # 归一化到 0-1
heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET)
overlay = cv2.addWeighted(image, 0.5, heatmap, 0.5, 0)
return overlay
def remove_module_prefix(state_dict):
new_state_dict = {}
for k, v in state_dict.items():
if k.startswith("module."):
new_state_dict[k[7:]] = v # 去掉前面的 'module.'
else:
new_state_dict[k] = v
return new_state_dict
# ========== 测试模型和 Grad-CAM ==========
if __name__ == "__main__":
# 1. 初始化模型（关键点数 16）
num_joints = 16
model = Your_model(num_joints, False)
# 2. 加载模型参数
state_dict = torch.load('model.pth')#你训练出的模型的路径
model.load_state_dict(remove_module_prefix(state_dict['model_state_dict']))
# 3. 选择目标层：绑定到模块中的某一层
target_layer = model.layer
# 4. 初始化 Grad-CAM
grad_cam = GradCAM(model, target_layer)
# 5. 读取测试图像并预处理
img = cv2.imread("test.jpg") # 读取图片（BGR 格式）
img = cv2.resize(img, (224, 224)) # 调整大小
# 转换为 RGB 格式，再转为 Tensor
img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img_tensor = torch.tensor(img_rgb.transpose(2, 0, 1), dtype=torch.float32).unsqueeze(0) / 255.0
# 6. 分别计算 0-16 号关键点的 Grad-CAM 热图（进行叠加）
cams = []
for joint_index in range(0,16):
cam_joint = grad_cam.forward(img_tensor, joint_index) # 每个 cam 的尺寸约为 40x40
# 将每个热图归一化到 0-1，并调整大小到原图尺寸
cam_joint = cv2.resize(cam_joint, (img.shape[1], img.shape[0]))
cam_joint = (cam_joint - cam_joint.min()) / (cam_joint.max() - cam_joint.min() + 1e-8)
cams.append(cam_joint)
cams = np.stack(cams, axis=0) # 形状 (16, H, W)
# 7. 合并 16 个关键点热图：逐像素取所有关键点中响应的最大值（不区分颜色，仅保留默认热图色调）
combined_cam = np.max(cams, axis=0)
combined_cam = (combined_cam - combined_cam.min()) / (combined_cam.max() - combined_cam.min() + 1e-8)
# 8. 生成叠加图（使用默认颜色映射）
overlay = grad_cam.overlay_heatmap(img, combined_cam)
plt.imshow(cv2.cvtColor(overlay, cv2.COLOR_BGR2RGB))
plt.axis("off")
plt.title("Combined Grad-CAM for Keypoints 0-20")
plt.show()