training scripts released

util/json_config.py

@@ -0,0 +1,15 @@
+import  json
+
+
+def readConfig(path):
+    with open(path,'r') as cf:
+        nodelocaltionstr = cf.read()
+        nodelocaltioninf = json.loads(nodelocaltionstr)
+        if isinstance(nodelocaltioninf,str):
+            nodelocaltioninf = json.loads(nodelocaltioninf)
+    return nodelocaltioninf
+
+def writeConfig(path, info):
+    with open(path, 'w') as cf:
+        configjson  = json.dumps(info, indent=4)
+        cf.writelines(configjson)

util/logo_class.py

@@ -0,0 +1,44 @@
+#!/usr/bin/env python3
+# -*- coding:utf-8 -*-
+#############################################################
+# File: logo_class.py
+# Created Date: Tuesday June 29th 2021
+# Author: Chen Xuanhong
+# Email: chenxuanhongzju@outlook.com
+# Last Modified:  Monday, 11th October 2021 12:39:55 am
+# Modified By: Chen Xuanhong
+# Copyright (c) 2021 Shanghai Jiao Tong University
+#############################################################
+
+class logo_class:
+    
+    @staticmethod
+    def print_group_logo():
+        logo_str = """
+
+███╗   ██╗██████╗ ███████╗██╗ ██████╗     ███████╗     ██╗████████╗██╗   ██╗
+████╗  ██║██╔══██╗██╔════╝██║██╔════╝     ██╔════╝     ██║╚══██╔══╝██║   ██║
+██╔██╗ ██║██████╔╝███████╗██║██║  ███╗    ███████╗     ██║   ██║   ██║   ██║
+██║╚██╗██║██╔══██╗╚════██║██║██║   ██║    ╚════██║██   ██║   ██║   ██║   ██║
+██║ ╚████║██║  ██║███████║██║╚██████╔╝    ███████║╚█████╔╝   ██║   ╚██████╔╝
+╚═╝  ╚═══╝╚═╝  ╚═╝╚══════╝╚═╝ ╚═════╝     ╚══════╝ ╚════╝    ╚═╝    ╚═════╝ 
+Neural Rendering Special Interesting Group of SJTU
+                                                                            
+        """
+        print(logo_str)
+
+    @staticmethod
+    def print_start_training():
+        logo_str = """
+   _____  __                __     ______              _         _              
+  / ___/ / /_ ____ _ _____ / /_   /_  __/_____ ____ _ (_)____   (_)____   ____ _
+  \__ \ / __// __ `// ___// __/    / /  / ___// __ `// // __ \ / // __ \ / __ `/
+ ___/ // /_ / /_/ // /   / /_     / /  / /   / /_/ // // / / // // / / // /_/ / 
+/____/ \__/ \__,_//_/    \__/    /_/  /_/    \__,_//_//_/ /_//_//_/ /_/ \__, /  
+                                                                       /____/   
+        """
+        print(logo_str)
+
+if __name__=="__main__":
+    # logo_class.print_group_logo()
+    logo_class.print_start_training()

util/plot.py

@@ -0,0 +1,37 @@
+import numpy as np
+import math
+import PIL
+
+def postprocess(x):
+    """[0,1] to uint8."""
+    
+    x = np.clip(255 * x, 0, 255)
+    x = np.cast[np.uint8](x)
+    return x
+
+def tile(X, rows, cols):
+    """Tile images for display."""
+    tiling = np.zeros((rows * X.shape[1], cols * X.shape[2], X.shape[3]), dtype = X.dtype)
+    for i in range(rows):
+        for j in range(cols):
+            idx = i * cols + j
+            if idx < X.shape[0]:
+                img = X[idx,...]
+                tiling[
+                        i*X.shape[1]:(i+1)*X.shape[1],
+                        j*X.shape[2]:(j+1)*X.shape[2],
+                        :] = img
+    return tiling
+
+
+def plot_batch(X, out_path):
+    """Save batch of images tiled."""
+    n_channels = X.shape[3]
+    if n_channels > 3:
+        X = X[:,:,:,np.random.choice(n_channels, size = 3)]
+    X = postprocess(X)
+    rc = math.sqrt(X.shape[0])
+    rows = cols = math.ceil(rc)
+    canvas = tile(X, rows, cols)
+    canvas = np.squeeze(canvas)
+    PIL.Image.fromarray(canvas).save(out_path)

util/save_heatmap.py

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