torch 显示注意力图

import torch
import numpy as np
import cv2
from PIL import Image
import matplotlib
#matplotlib.use('Agg') #注释掉则不show图
import matplotlib.pyplot as plt
import os


def visulize_attention_ratio(img_path, save_path, attention_mask, ratio=0.5, cmap="jet", save_image=False,
                             save_original_image=False):
    """
    img_path:   image file path to load
    save_path:  image file path to save
    attention_mask: 2-D attention map with np.array type, e.g, (h, w) or (w, h)
    ratio:  scaling factor to scale the output h and w
    cmap:   attention style, default: "jet"
    """
    print("load image from: ", img_path)
    img = Image.open(img_path, mode='r')
    img_h, img_w = img.size[0], img.size[1]
    plt.subplots(nrows=1, ncols=1, figsize=(0.02 * img_h, 0.02 * img_w))

    # scale the image
    img_h, img_w = int(img_h * ratio), int(img_w * ratio)
    img = img.resize((img_h, img_w))
    plt.imshow(img, alpha=1)
    plt.axis('off')

    mask = cv2.resize(attention_mask, (img_h, img_w))
    normed_mask = mask / mask.max()
    normed_mask = (normed_mask * 255).astype('uint8')
    plt.imshow(normed_mask, alpha=0.5, interpolation='nearest', cmap=cmap)

    if save_image:
        img_name = img_path.split('/')[-1].split('.')[0] + "_with_attention.jpg"
        img_with_attention_save_path = os.path.join(save_path, img_name)
        # pre-process before saving
        print("save image to: " + save_path)
        plt.axis('off')
        plt.subplots_adjust(top=1, bottom=0, right=1, left=0, hspace=0, wspace=0)
        plt.margins(0, 0)
        plt.savefig(img_with_attention_save_path, dpi=100)

    if save_original_image:
        print("save original image at the same time")
        img_name = img_path.split('/')[-1].split('.')[0] + "_original.jpg"
        original_image_save_path = os.path.join(save_path, img_name)
        img.save(original_image_save_path, quality=100)
    


from torchvision import transforms

img = Image.open('original_fig.jpg', mode='r')

transform = transforms.Compose([ transforms.ToTensor()])  
tensor = transform(img)

visulize_attention_ratio(img_path='original_fig.jpg', 
                         save_path='fig1.jpg', 
                         attention_mask=tensor[0,:,:].numpy(),
                         ratio=0.5, cmap="jet", save_image=True, save_original_image=False)