MMSegmentation绘制热力图

在MMSegmentation的一些版本例如v1.2.0中，官方已经提供了利用Grad-CAM绘制热力图的代码，但只是针对单张图片，对于论文写作的朋友们来说不太友好，笔者将其修改为能够批量处理文件夹从而提高绘制效率，同时，针对里面一些不太容易弄懂的地方笔者用注释写出了自己的理解。修改代码如下：

"""
Use the pytorch-grad-cam tool to visualize Class Activation Maps (CAM).

requirement: pip install grad-cam
"""

from argparse import ArgumentParser
import os

import numpy as np
import torch
import torch.nn.functional as F
from mmengine import Config
from mmengine.model import revert_sync_batchnorm
from PIL import Image
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.image import preprocess_image, show_cam_on_image

from mmseg.apis import inference_model, init_model, show_result_pyplot
from mmseg.utils import register_all_modules


class SemanticSegmentationTarget:
    """wrap the model.

    requirement: pip install grad-cam

    Args:
        category (int): Visualization class.
        mask (ndarray): Mask of class.
        size (tuple): Image size.
    """

    def __init__(self, category, mask, size):
        self.category = category
        self.mask = torch.from_numpy(mask)
        self.size = size
        if torch.cuda.is_available():
            self.mask = self.mask.cuda()

    def __call__(self, model_output):
        model_output = torch.unsqueeze(model_output, dim=0)
        model_output = F.interpolate(
            model_output, size=self.size, mode='bilinear', align_corners=True)
        model_output = torch.squeeze(model_output, dim=0)

        return (model_output[self.category, :, :] * self.mask).sum()


def main():
    parser = ArgumentParser()

    #设置预测文件夹路径
    parser.add_argument('--input_dir',default='./img_predict/imgs512', help='Input directory containing images')
    
    #设置配置文件路径
    parser.add_argument('--config', default='./my_moudel/deeplabv3plus/deeplabv3plus_r50-d8_4xb4-40k_voc12aug-512x512.py',help='Config file')

    #设置权重文件路径
    parser.add_argument('--checkpoint', default='./work_dirs/rccnet_cbam_cat/best_mIoU_iter_9550.pth', help='Checkpoint file')

    #设置保存输出的热力图文件夹路径
    parser.add_argument(
        '--out-dir',
        default='./img_predict/imgs512_cam_1',
        help='Directory to save CAM images')

    #设置需要输出网络特定层的热力图
    parser.add_argument(
        '--target-layers',
        default='decode_head.bottleneck.activate',
        help='Target layers to visualize CAM')

    #Grad-CAM为类激活加权映射，所以需要设置固定类别的索引来绘制，需要与传入网络训练时候一致
    parser.add_argument(
        '--category-index', default='1', help='Category to visualize CAM')
    parser.add_argument(
        '--device', default='cuda:0', help='Device used for inference')
    args = parser.parse_args()

    # 创建保存目录
    os.makedirs(args.out_dir, exist_ok=True)

    # build the model from a config file and a checkpoint file
    register_all_modules()
    model = init_model(args.config, args.checkpoint, device=args.device)
    if args.device == 'cpu':
        model = revert_sync_batchnorm(model)

    # 列出输入目录中的所有图片文件
    input_files = [f for f in os.listdir(args.input_dir) if f.endswith(('.jpg', '.jpeg', '.png'))]

    for input_file in input_files:
        input_path = os.path.join(args.input_dir, input_file)

        # test a single image
        result = inference_model(model, input_path)

        # result data conversion
        prediction_data = result.pred_sem_seg.data
        pre_np_data = prediction_data.cpu().numpy().squeeze(0)

        target_layers = args.target_layers
        target_layers = [eval(f'model.{target_layers}')]

        category = int(args.category_index)
        mask_float = np.float32(pre_np_data == category)

        # data processing
        image = np.array(Image.open(input_path).convert('RGB'))
        height, width = image.shape[0], image.shape[1]
        rgb_img = np.float32(image) / 255
        config = Config.fromfile(args.config)
        image_mean = config.data_preprocessor['mean']
        image_std = config.data_preprocessor['std']
        input_tensor = preprocess_image(
            rgb_img,
            mean=[x / 255 for x in image_mean],
            std=[x / 255 for x in image_std])

        # Grad CAM(Class Activation Maps)
        
        targets = [
            SemanticSegmentationTarget(category, mask_float, (height, width))
        ]
        cam_file_name = os.path.splitext(input_file)[0] + '_cam.jpg'  # 使用输入图片文件名构造保存文件名
        cam_file_path = os.path.join(args.out_dir, cam_file_name)
        # Can also be LayerCAM, XGradCAM, GradCAMPlusPlus, EigenCAM, EigenGradCAM，这里可以根据需要选择不同的特征图生成方式
        with GradCAM(
                model=model,
                target_layers=target_layers,
#                 use_cuda=torch.cuda.is_available()
        ) as cam:
            grayscale_cam = cam(input_tensor=input_tensor, targets=targets)[0, :]
            cam_image = show_cam_on_image(rgb_img, grayscale_cam, use_rgb=True)

            # 保存 CAM 图片
            Image.fromarray(cam_image).save(cam_file_path)


if __name__ == '__main__':
main()

下面给出了上述提到的多种热力图绘制方法的区别，笔者个人还是倾向使用GradCAM： 

1. GradCAM (Gradient-weighted Class Activation Mapping):

   • 原理：GradCAM 利用最后的卷积层的梯度信息来生成类别激活图。具体来说，它通过计算目标类别相对于卷积层输出的梯度，得到每个通道的重要性权重，然后使用这些权重对卷积层的特征图进行加权求和，最终生成一个热力图来表示图像中对该类别最重要的区域。
   • 优点：直观简单，能够很好地捕捉到目标对象的整体轮廓。
   • 缺点：由于是基于卷积层输出的平均加权，有时可能会忽略局部细节。

2. LayerCAM:

   • 原理：LayerCAM 是 GradCAM 的变体，通过逐层逐通道地计算激活图，而不是在最后一层进行汇总。这样可以捕捉到更多的细节信息。
   • 优点：更细致地展示目标区域，适合需要更高分辨率的情况。
   • 缺点：计算量相对较大。

3. XGradCAM:

   • 原理：XGradCAM 是 GradCAM 的扩展，结合了不同层的梯度信息。它通过跨层级的梯度信息，增强了类别激活图的表达能力。
   • 优点：能够更全面地捕捉目标对象的特征，综合多层信息，生成的热力图更加精细。
   • 缺点：由于结合了多层信息，计算开销较大。

4. GradCAMPlusPlus:

   • 原理：GradCAMPlusPlus 通过对 GradCAM 的改进，更加准确地捕捉图像中每个像素对于类别的贡献。它利用了更复杂的权重计算方法，考虑了不同通道之间的差异。
   • 优点：在处理多目标场景时更加有效，生成的激活图更为精细和准确。
   • 缺点：计算复杂度较高。

5. EigenCAM:

   • 原理：EigenCAM 利用卷积层的特征图的主成分分析 (PCA) 来生成激活图。具体来说，它通过计算特征图的前几个主成分，生成反映目标类别的激活图。
   • 优点：能够有效地提取主要特征，计算效率较高。
   • 缺点：对于一些复杂场景，可能不如 GradCAM 和 GradCAMPlusPlus 精细。

6. EigenGradCAM:

   • 原理：EigenGradCAM 结合了 GradCAM 和 EigenCAM 的方法，既利用了梯度信息，又利用了主成分分析，生成更加准确的激活图。
   • 优点：综合了两种方法的优点，生成的激活图更加全面和准确。
   • 缺点：计算复杂度较高，需要更多的计算资源。

选择合适的 CAM 方法

• GradCAM：适合一般场景，计算效率高，能够很好地捕捉目标的整体轮廓。
• LayerCAM：适合需要高分辨率和细节的应用场景，计算量相对较大。
• XGradCAM：适合需要综合多层信息的场景，生成的激活图更加全面。
• GradCAMPlusPlus：适合多目标场景，生成的激活图更为精细和准确。
• EigenCAM：适合需要高效计算的场景，利用主成分分析提取主要特征。
• EigenGradCAM：综合了 GradCAM 和 EigenCAM 的优点，适合需要更加全面和准确的激活图的场景。