init

2022-01-10 15:03:58 +08:00
parent 573689a591
commit 3783ef0e75
57 changed files with 9520 additions and 0 deletions
@@ -0,0 +1,57 @@
+#!/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:  Friday, 15th January 2021 10:23:13 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)