|
|
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] |
|
|
return attention_img |
