File size: 922 Bytes
e34aada |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 |
import numpy as np
import cv2
from utils.commons.tensor_utils import convert_to_np
def plot_attention_img(attention_img, color_bar='jet'):
"""
attention_img: raw attention in network, tensor or array, in 0~1 scale, shape [H, W,]
color_bar: jet, summer, etc see this https://blog.csdn.net/loveliuzz/article/details/73648505
return: ready-to-visualize attention img in -1~1 scale.
"""
attention_img = convert_to_np(attention_img)
assert attention_img.ndim == 2
attention_img = np.uint8(255 * attention_img)
color_bar_dict = {
'jet': cv2.COLORMAP_JET,
'summer': cv2.COLORMAP_SUMMER,
'hot': cv2.COLORMAP_HOT
}
color_bar = color_bar_dict.get(color_bar, getattr(cv2, f"COLORMAP_{color_bar.upper()}"))
attention_img = cv2.applyColorMap(attention_img, color_bar) / 127.5 - 1
attention_img = attention_img[:, :, ::-1] # flip RGB
return attention_img |