|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
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) |
|
|
|
