vis_feature_map.py

from argparse import ArgumentParser
from typing import Type

import mmcv
import torch
import torch.nn as nn
import numpy as np

from mmengine.model import revert_sync_batchnorm
from mmengine.structures import PixelData
from mmseg.apis import inference_model, init_model
from mmseg.structures import SegDataSample
from mmseg.utils import register_all_modules
from mmseg.visualization import SegLocalVisualizer

from src.apis import inference_model


class Recorder:
    """record the forward output feature map and save to data_buffer."""

    def __init__(self) -> None:
        self.data_buffer = list()

    def __enter__(self, ):
        self._data_buffer = list()

    def record_data_hook(self, model: nn.Module, input: Type, output: Type):
        self.data_buffer.append(output)

    def __exit__(self, *args, **kwargs):
        pass


def visualize(args, model, recorder, result):
    seg_visualizer = SegLocalVisualizer(
        vis_backends=[
            # dict(type='WandbVisBackend'),
            dict(type='LocalVisBackend'),
        ],
        save_dir='UperNet_Featmap',
        alpha=0.5)
    seg_visualizer.dataset_meta = dict(classes=model.dataset_meta['classes'],
                                       palette=model.dataset_meta['palette'])

    image = mmcv.imread(args.img1, 'color')
    image2 = mmcv.imread(args.img2, 'color')

    seg_visualizer.add_datasample(name='predict',
                                  image=image,
                                  data_sample=result,
                                  draw_gt=False,
                                  draw_pred=True,
                                  wait_time=0,
                                  out_file=None,
                                  show=False,
                                  withLabels=False)

    # add feature map to wandb visualizer
    for i in range(len(recorder.data_buffer)):
        rec = recorder.data_buffer[i]
        if isinstance(rec, torch.Tensor):
            feature = rec[0]  # remove the batch
            drawn_img = seg_visualizer.draw_featmap(
                feature, image, channel_reduction='select_max')
            seg_visualizer.add_image(f'feature_map{i}', drawn_img)
        elif isinstance(rec, (list, tuple)):
            for j, f in enumerate(rec):
                feature = f[0]  # remove the batch
                drawn_img = seg_visualizer.draw_featmap(
                    feature, image, channel_reduction='select_max')
                seg_visualizer.add_image(f'feature_map{i}_{j}', drawn_img)

    if args.gt_mask:
        sem_seg = mmcv.imread(args.gt_mask, 'unchanged')
        sem_seg -= 1
        sem_seg = torch.from_numpy(sem_seg)
        gt_mask = dict(data=sem_seg)
        gt_mask = PixelData(**gt_mask)
        data_sample = SegDataSample()
        data_sample.gt_sem_seg = gt_mask

        seg_visualizer.add_datasample(name='gt_mask',
                                      image=image,
                                      data_sample=data_sample,
                                      draw_gt=True,
                                      draw_pred=False,
                                      wait_time=0,
                                      out_file=None,
                                      show=False,
                                      withLabels=False)

    seg_visualizer.add_image('image', image)
    seg_visualizer.add_image('image2', image2)


def main():
    parser = ArgumentParser(
        description='Draw the Feature Map During Inference')
    parser.add_argument('img1', help='Image file')
    parser.add_argument('img2', help='Image file')
    parser.add_argument('config', help='Config file')
    parser.add_argument('checkpoint', help='Checkpoint file')
    parser.add_argument('--gt_mask', default=None, help='Path of gt mask file')
    parser.add_argument('--out-file', default=None, help='Path to output file')
    parser.add_argument('--device',
                        default='cuda:0',
                        help='Device used for inference')
    parser.add_argument(
        '--opacity',
        type=float,
        default=0.5,
        help='Opacity of painted segmentation map. In (0, 1] range.')
    parser.add_argument('--title',
                        default='result',
                        help='The image identifier.')
    args = parser.parse_args()

    register_all_modules()

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

    # show all named module in the model and use it in source list below
    for name, module in model.named_modules():
        print(name)

    source = [
        # UperNet
        'backbone.backbone.stages.3',
        'backbone.backbone2.stages.3',
        # # ConvNeXt
        # 'backbone.backbone.stages.3',
        # 'backbone.backbone2.stages.3',
        # # 'backbone.backbone.stages.2',
        # # 'backbone.backbone.stages.0',
        # 'backbone.FRMs.3',
        # 'backbone.FFMs.3',
        # 'decode_head.fpn_bottleneck',
        # # CMX
        # 'backbone.block3_out',
        # 'backbone.extra_block3_out',
        # 'backbone.FRMs.2',
        # 'backbone.FFMs.2',
    ]
    source = dict.fromkeys(source)

    count = 0
    recorder = Recorder()
    # registry the forward hook
    for name, module in model.named_modules():
        if name in source:
            count += 1
            module.register_forward_hook(recorder.record_data_hook)
            if count == len(source):
                break

    with recorder:
        # test a single image, and record feature map to data_buffer
        result = inference_model(model, args.img1, args.img2)

    visualize(args, model, recorder, result)


if __name__ == '__main__':
    main()