feature map 可视化之目标检测（更新）

看了下CAM（class activation mapping）相关文章，想着自己搞一下YOLOV5 最后输出的三层head的feature map可视化，好像有点效果。

输出的三张热力图是：

第一张: 只靠confidence进行feature map映射

第二张: 只靠classes 概率进行feature map映射

第三张: 将confidence*classes 进行feature map映射

第一层head效果：

第二层head效果：

第三层head效果：

整合3层head效果：

修改：2021-8-18

最近工作需要重新整理下这个代码，方便大家使用，案例采用YOLOV5为例。

#从yolov5官方代码的detect.py中找到如下代码：
#测试过程中batchsize必须设置为1

pred = model(img, augment=opt.augment)[0]
#修改为

ret= model(img, augment=opt.augment)
pred = ret[0]

#再在后面添加如下代码
# 下面的1指的是你想查看的类别；80+5指的是coco数据集的80个类别+ x y w h score 5个数值
#    image_size：模型输入尺寸
#    normalization：如果为True将置信度和class归一化到0~1，方便按置信度进行区分热力图，如果为False会按本身数据分布归一化，这样方便查看相对置信度。
# image: 原图 即 cv2.imread() 的结果
    
show_CAM("热力图保存文件夹", image, ret[1], 1, 80+5， image_size=(320, 640), normalization=True)

一下是代码：

import os

import cv2
import numpy as np
import torch


def show_CAM(save_img_path, image, feature_maps, class_id, all_ids=10, image_size=(320, 640), normalization=True):
    """
    save_img_path: save heatmap images path
    feature_maps: this is a list [tensor,tensor,tensor], tensor shape is [1, 3, N, N, all_ids]
    normalization: Normalize score and class to 0 to 1
    image_size: w, h
    """
    SHOW_NAME = ["score", "class", "class*score"]
    img_ori = image
    layers0 = feature_maps[0].reshape([-1, all_ids])
    layers1 = feature_maps[1].reshape([-1, all_ids])
    layers2 = feature_maps[2].reshape([-1, all_ids])
    layers = torch.cat([layers0, layers1, layers2], 0)
    if normalization:
        score_max_v = 1.
        score_min_v = 0.
        class_max_v = 1.
        class_min_v = 0.
    else:
        score_max_v = layers[:, 4].max()  # compute max of score from all anchor
        score_min_v = layers[:, 4].min()  # compute min of score from all anchor
        class_max_v = layers[:, 5 + class_id].max()  # compute max of class from all anchor
        class_min_v = layers[:, 5 + class_id].min()  # compute min of class from all anchor
    for j in range(3):  # layers
        layer_one = feature_maps[j]
        # compute max of score from three anchor of the layer
        if normalization:
            anchors_score_max = layer_one[0, :, :, :, 4].max(0)[0].sigmoid()
            # compute max of class from three anchor of the layer
            anchors_class_max = layer_one[0, :, :, :, 5 + class_id].max(0)[0].sigmoid()
        else:
            anchors_score_max = layer_one[0, :, :, :, 4].max(0)[0]
            # compute max of class from three anchor of the layer
            anchors_class_max = layer_one[0, :, :, :, 5 + class_id].max(0)[0]

        scores = ((anchors_score_max - score_min_v) / (
                score_max_v - score_min_v))
        classes = ((anchors_class_max - class_min_v) / (
                class_max_v - class_min_v))

        layer_one_list = []
        layer_one_list.append(scores)
        layer_one_list.append(classes)
        layer_one_list.append(scores * classes)
        for idx, one in enumerate(layer_one_list):
            layer_one = one.cpu().numpy()
            if normalization:
                ret = ((layer_one - layer_one.min()) / (layer_one.max() - layer_one.min())) * 255
            else:
                ret = ((layer_one - 0.) / (1. - 0.)) * 255
            ret = ret.astype(np.uint8)
            gray = ret[:, :, None]
            ret = cv2.applyColorMap(gray, cv2.COLORMAP_JET)

            ret = cv2.resize(ret, image_size)
            img_ori = cv2.resize(img_ori, image_size)

            show = ret * 0.50 + img_ori * 0.50
            show = show.astype(np.uint8)
            cv2.imwrite(os.path.join(save_img_path, f"{j}_{SHOW_NAME[idx]}.jpg"), show)


# show_CAM(path, ret[1], 21)

自己瞎捣鼓的，有错误的地方大家提出来啊~

版权声明：本文为CSDN博主「One橙序猿」的原创文章，遵循CC 4.0 BY-SA版权协议，转载请附上原文出处链接及本声明。
原文链接：https://blog.csdn.net/rush9838465/article/details/114291144