[PyTorch]即插即用的热力图生成

先上张效果图,本来打算移植霹雳老师的使用Pytorch实现Grad-CAM并绘制热力图。但是看了下代码,需要骨干网络按照标准写法(即将特征层封装为features数组),而我写的网络图省事并没有进行封装,改造网络的代价又太大了,所以干脆直接重写一个。

一、生成热力图

大致可以分为三步:①读取图片;②前向传递运算;③用特征向量生成特征图。而图片的resize图简单可以直接用t**ransforms,**后面反正也是直接resize回来的,并不会造成变形。

python 复制代码
# 加载一个transforms用于变形,input_shape为预设的图像尺寸
transform = transforms.Compose([transforms.Resize((input_shape[0],input_shape[1])),
                                    transforms.ToTensor(),
                                    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),])
image = Image.open(image_path)     #image_path为文件路径
input_tensor = transform(image)    #将图片转换为tensor类型
input_batch = input_tensor.unsqueeze(0)    #为tensor添加batch维度

# 前向传递
model.eval()
with torch.no_grad():
    output = model(input_batch)

使用特征图生成热力图的原理是:将该维度上所有的tensor进行叠加,然后将生成的矩阵变形回输入向量的尺寸

python 复制代码
heatmap = torch.sum(output, dim=1)    #所有通道求和
max_value = torch.max(heatmap)
min_value = torch.min(heatmap)
heatmap = (heatmap-min_value)/(max_value-min_value)*255

heatmap = heatmap.cpu().numpy().astype(np.uint8).transpose(1,2,0)  # 提取热力图

heatmap = cv2.resize(heatmap, input_shape,interpolation=cv2.INTER_LINEAR)  # 还原尺寸

# 将矩阵转换为image类
heatmap=cv2.applyColorMap(heatmap,cv2.COLORMAP_JET)
heatimg = Image.fromarray(heatmap)

二、叠加原图

直接使用plt进行叠加!

python 复制代码
    # 将热力图叠加到原图上
    org_size = image.size
    heatimg = heatimg.resize(org_size)    #将热力图变回输入图像的尺寸
    plt.axis('off')
    plt.imshow(image)
    plt.imshow(heatimg, alpha=0.5)  # alpha为热力图的透明度

    # 显示叠加后的图形
    plt.show()

三、总结

这段代码和霹雳老师的Grad-CAM对比优劣都很明显,优点是代码比较简单。上可以通过插入前向传递的环境直接得到任何层的热力图。但缺点就是不能关注特定的类别,且生成的热力图也不是很美观。

相关推荐
whaosoft-14335 分钟前
51c自动驾驶~合集7
人工智能
都叫我大帅哥40 分钟前
Python的Optional:让你的代码优雅处理“空值”危机
python
曾几何时`3 小时前
基于python和neo4j构建知识图谱医药问答系统
python·知识图谱·neo4j
刘晓倩4 小时前
Coze智能体开发实战-多Agent综合实战
人工智能·coze
石迹耿千秋5 小时前
迁移学习--基于torchvision中VGG16模型的实战
人工智能·pytorch·机器学习·迁移学习
写写闲篇儿5 小时前
Python+MongoDB高效开发组合
linux·python·mongodb
杭州杭州杭州6 小时前
Python笔记
开发语言·笔记·python
路人蛃8 小时前
通过国内扣子(Coze)搭建智能体并接入discord机器人
人工智能·python·ubuntu·ai·aigc·个人开发
CV-杨帆8 小时前
论文阅读:arxiv 2025 A Survey of Large Language Model Agents for Question Answering
论文阅读·人工智能·语言模型
qiqiqi(^_×)8 小时前
卡在“pycharm正在创建帮助程序目录”
ide·python·pycharm