import matplotlib.pyplot as plt
import cv2
f = pred_score.squeeze().cpu().detach().numpy()
f = f.transpose(2,1,0)[:,:,9]
# print(f.shape)#(17,17)
# a = np.stack([cv2.resize(f,(17,17),interpolation=cv2.INTER_CUBIC)])#a.shape->(1, 17, 17)
# a = np.squeeze(a)
# # a = np.mean(f)
a = np.maximum(f,0)
a /= np.max(a)
heatmap = cv2.resize(a, (img.shape[1], img.shape[0]))#img.shape->(255,255,3)
heatmap = np.uint8(255 * heatmap)
heatmap = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET)
heatmap = heatmap * 0.5 + img
cv2.imwrite('保存的图片路径', heatmap)
效果图
opencv中伪彩色applyColorMap函数转发自下面链接
https://blog.csdn.net/u013381011/article/details/78341861