pytorch FPN层objectness map的可视化

要解决的问题

FPN Faster R-CNN的多尺度特征包括P2, P3, P4, P5, P6，目的是观察RPN网络所生成的预测框的前景得分热力图，以验证FPN的不同层确实将不同尺度的特征进行了融合，下面是所得到的的示例图，因为我本身是研究细胞图像的，下面就以细胞图像为例：

所采用的方法

FPN Faster R-CNN结合score cam进行热力图的呈现

1. pytorch FPN Faster R-CNN模型的构建和权重加载，这里的权重是我使用自己的数据集进行训练得到的，我本身对源码进行了修改适应我的研究，我采用的是densenet161，所以和官方代码稍有不同，因为官方代码目前支持的是resnet系列，所以下面的代码是以官方代码为例，具体情况需要自己进行修改以适应自己的研究。

# 导入pytorch官方的Faster R-CNN相关的模块
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor, FastRCNN
from torchvision.models.detection.backbone_utils import resnet_fpn_backbone
import torch
import numpy as np
from PIL import Image
from torchvision import transforms
import matplotlib.cm

import copy

# 特征提取的骨架网络外加FPN模块构成backbone
# 这里的resnet50使用ImageNet训练得到的权重
backbone = resnet_fpn_backbone("resnet50", True)
# 很简单地就将FPN Faster R-CNN模型构建好了
model = FasterRCNN(backbone, num_classes=2)  # 我的类别是2类

2. 重头戏就是该网络RPN部分cls_logits的输出，这属于中间层特征的输出，所以需要用到pytorch的hook功能，这是pytorch官方提出的一个很方便的用于进行中间层梯度，特征等提取的功能，具体详细的介绍以后再说，这里只是简单地利用一下，非常简单易于操作。

# 定义列表用于存储中间层的输入或者输出
p_in = []
p_out = []

# 定义hook_fn
def hook_fn(module, inputs, outputs):
    p_in.append(inputs)
    p_out.append(outputs)
# 很随意的就将中间层结果存储起来了，需要将hook_fn“挂”到“目标层”上
model.rpn.cls_logits.register_forward_hook(hook_fn)
# 这样就OK了，进行前向运算的时候，hook_fn就自动的执行

3. 前向运算，并将获取的cls_logits转换为热力图覆盖到原图像

# 导入一张图像
img = Image.open("xxx.png")
ori_img = img.copy()

# 必要的前处理
transform = transforms.Compose([
				transforms.ToTensor(),
				transforms.Normalize(
					[0.485, 0.456, 0.406],
					[0.229, 0.224, 0.225])])
img = transform(img).unsqueeze(0)  # 增加batch维度
# 首先model需要处于eval模式
model.eval()
with torch.no_grad():
    model(img)

前向传播结束以后，cls_logits就已经存储到了p_out列表中，总共有5个元素，分别代表P2, P3, P4, P5, P6特征图锚点框所对应的预测框前景类别得分logits值

4. 得到cls_logits值以后，需要对其进行处理，使其变为0-255的热力图，需要注意的是每个水平的特征所对应的cls_logits都是一个3通道的张量（该数值取决于anchor_scale和aspect_ratio的取值，我取的是每个锚点处3种不同高宽比的框，所以是3通道张量），为了展示热力图，我们需要从3通道中取出一个通道，也就是一种高宽比的锚点框。

# 以P2水平为例，首先将cls_logits经过sigmoid函数得到0-1之间的值
p2_logits = torch.sigmoid(p_out[0])
p2_logits = p2_logits.detach().numpy()  # 转为numpy数组
# 去除小于0的值，该做法参考的是score cam文章的做法
p2_logits = np.maximum(p2_logits, 0)
# 归一化
p2_logits = (p2_logits - np.min(p2_logits))/(np.max(p2_logits)-np.min(p2_logits))
# 转为uint8类型
p2_logits = np.uint8(p2_logits * 255)
# 取其中一个通道
p2_logits = p2_logits[0, 0, ...]
# 缩放至原图大小，我的原图大小为2048*2048
p2_logits = np.uint8(Image.fromarray(p2_logits).resize((2048, 2048), Image.ANTIALIAS)) / 255

这样一来，热力图就得到了，但是还需将其覆盖到原图

5. 定义函数进行覆盖操作

def put_heatmap_on_image(ori_image, activation, colormap_name):
    """
    ori_image (PIL image): 原始图像
    activation (numpy arr): 即上面得到的p2_logits
    colormap_name (str): 采用何种matplotlib.cm的colormap
    """
    # colormap
    color_map = matplotlib.cm.get_map(colormap_name)
    # 把colormap添加到activation，即activation的以何种
    # colormap进行显示
    no_trans_heatmap = color_map(activation)
    # 添加alpha通道，即透明度
    heatmap = copy.copy(no_trans_heatmap)
    heatmap[:, :, 3] = 0.4
    heatmap = Image.fromarray((heatmap*255).astype(np.uint8))
    no_trans_heatmap = Image.fromarray((no_trans_heatmap*255).astype(np.uint8)
    # 把热力图覆盖到原图像
    heatmap_on_image = Image.new("RGBA", ori_image.size)
    heatmap_on_image = Image.alpha_composite(
    					heatmap_on_image, ori_image.convert("RGBA"))
    heatmap_on_image = Image.alpha_composite(
    					heatmap_on_image, heatmap)
    return no_trans_heatmap, heatmap_on_image


heatmap, heatmap_on_image = put_heatmap_on_image(
			ori_image, p2_logits, "jet")
heatmap_on_image.save("xxx.png")

这样一来，就将最终的结果图像存储起来了