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")

這樣一來，就將最終的結果圖像存儲起來了