基於mmdetection的熱力圖可視化繪制

圖例(隨便找的一張圖，訓練效果比較糟糕)：

 

用例：

/home/aimhabo/anaconda3/envs/mmlab/bin/python ./tools/visualization.py 0000308_04401_d_0000327.jpg ./featuremaps/ retina_res50_fpn.py retina_res50_fpn.pth

 

visualization.py：

#coding: utf-8
import cv2
import mmcv
import numpy as np
import os
import torch

from mmdet.apis import inference_detector, init_detector

def featuremap_2_heatmap(feature_map):
    assert isinstance(feature_map, torch.Tensor)
    feature_map = feature_map.detach()
    heatmap = feature_map[:,0,:,:]*0
    for c in range(feature_map.shape[1]):
        heatmap+=feature_map[:,c,:,:]
    heatmap = heatmap.cpu().numpy()
    heatmap = np.mean(heatmap, axis=0)

    heatmap = np.maximum(heatmap, 0)
    heatmap /= np.max(heatmap)

    return heatmap

def draw_feature_map(model, img_path, save_dir):
    '''
    :param model: 加載了參數的模型
    :param img_path: 測試圖像的文件路徑
    :param save_dir: 保存生成圖像的文件夾
    :return:
    '''
    img = mmcv.imread(img_path)
    modeltype = str(type(model)).split('.')[-1].split('\'')[0]
    model.eval()
    model.draw_heatmap = True
    featuremaps = inference_detector(model, img) #1.這里需要改model，讓其在forward的最后return特征圖。我這里return的是一個Tensor的tuple，每個Tensor對應一個level上輸出的特征圖。
    i=0
    for featuremap in featuremaps:
        heatmap = featuremap_2_heatmap(featuremap)
        heatmap = cv2.resize(heatmap, (img.shape[1], img.shape[0]))  # 將熱力圖的大小調整為與原始圖像相同
        heatmap = np.uint8(255 * heatmap)  # 將熱力圖轉換為RGB格式
        heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)  # 將熱力圖應用於原始圖像
        superimposed_img = heatmap * 0.4 + img  # 這里的0.4是熱力圖強度因子
        cv2.imwrite(os.path.join(save_dir,'featuremap_'+str(i)+'.png'), superimposed_img)  # 將圖像保存到硬盤
        i=i+1


from argparse import ArgumentParser

def main():
    parser = ArgumentParser()
    parser.add_argument('img', help='Image file')
    parser.add_argument('save_dir', help='Dir to save heatmap image')
    parser.add_argument('config', help='Config file')
    parser.add_argument('checkpoint', help='Checkpoint file')
    parser.add_argument('--device', default='cuda:0', help='Device used for inference')
    args = parser.parse_args()

    # build the model from a config file and a checkpoint file
    model = init_detector(args.config, args.checkpoint, device=args.device)
    draw_feature_map(model,args.img,args.save_dir)

if __name__ == '__main__':
    main()

 

對於其中1. 所述的model修改，以two_stage及其衍生方法為例，是指在simple_test和aug_test中，在extract_feats之后直接return (x,)