合肥电子商务开发网站建设手机网站快速建站
不好之处就是这种方法不能最大程度还原最后的热图,会产生很多噪声,不过大体区域还是接近的,代码如下:
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 = modelself.target_layer = target_layerself.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 camdef 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-1heatmap = cv2.applyColorMap(np.uint8(255 * cam), cv2.COLORMAP_JET)overlay = cv2.addWeighted(image, 0.5, heatmap, 0.5, 0)return overlaydef 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] = vreturn new_state_dict# ========== 测试模型和 Grad-CAM ==========
if __name__ == "__main__":# 1. 初始化模型(关键点数 16)num_joints = 16model = 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-CAMgrad_cam = GradCAM(model, target_layer)# 5. 读取测试图像并预处理img = cv2.imread("test.jpg") # 读取图片(BGR 格式)img = cv2.resize(img, (224, 224)) # 调整大小# 转换为 RGB 格式,再转为 Tensorimg_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()