File size: 1,702 Bytes
f884940 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
#!/usr/bin/env python3
# -*- coding:utf-8 -*-
#############################################################
# File: save_heatmap.py
# Created Date: Friday January 15th 2021
# Author: Chen Xuanhong
# Email: chenxuanhongzju@outlook.com
# Last Modified: Wednesday, 19th January 2022 1:22:47 am
# Modified By: Chen Xuanhong
# Copyright (c) 2021 Shanghai Jiao Tong University
#############################################################
import os
import shutil
import seaborn as sns
import matplotlib.pyplot as plt
import cv2
import numpy as np
def SaveHeatmap(heatmaps, path, row=-1, dpi=72):
"""
The input tensor must be B X 1 X H X W
"""
batch_size = heatmaps.shape[0]
temp_path = ".temp/"
if not os.path.exists(temp_path):
os.makedirs(temp_path)
final_img = None
if row < 1:
col = batch_size
row = 1
else:
col = batch_size // row
if row * col <batch_size:
col +=1
row_i = 0
col_i = 0
for i in range(batch_size):
img_path = os.path.join(temp_path,'temp_batch_{}.png'.format(i))
sns.heatmap(heatmaps[i,0,:,:],vmin=0,vmax=heatmaps[i,0,:,:].max(),cbar=False)
plt.savefig(img_path, dpi=dpi, bbox_inches = 'tight', pad_inches = 0)
img = cv2.imread(img_path)
if i == 0:
H,W,C = img.shape
final_img = np.zeros((H*row,W*col,C))
final_img[H*row_i:H*(row_i+1),W*col_i:W*(col_i+1),:] = img
col_i += 1
if col_i >= col:
col_i = 0
row_i += 1
cv2.imwrite(path,final_img)
if __name__ == "__main__":
random_map = np.random.randn(16,1,10,10)
SaveHeatmap(random_map,"./wocao.png",1)
|