change functions to plot image

libs/plotting.py

@@ -107,4 +107,34 @@ def plotRSS(
     plt.close()
 
 
+def plotSingleImage(
+    data: np.ndarray,
+    display: bool = False,
+    dump: bool = False,
+    fp: str = None,
+    cminmax: tuple = (COLORMAP_MIN, COLORMAP_MAX),
+    title: str = "map"
+):
+    '''
+    '''
+    plt.figure(figsize=(4, 4))
+    plt.title(title)
+    plt.imshow(
+        np.transpose(data),
+        cmap='hot',
+        origin='lower',
+        interpolation='nearest',
+        vmin=cminmax[0],
+        vmax=cminmax[1]
+    )
+    plt.colorbar()
+
+    plt.draw()
+    if dump:
+        plt.savefig(fp.replace('.pickle', '.png'), dpi=50)
+    if display:
+        plt.pause(0.1)
+        q = input("press enter..")
+    plt.close()
+
 

visualize.py

@@ -5,7 +5,7 @@ import sys
 import pickle
 import argparse
 
-from libs.plotting import plotRSS
+from libs.plotting import plotSingleImage
 
 
 def loadData(filepath):
@@ -25,7 +25,17 @@ def plotSingle(filepath, args):
     if data is None:
         print("err: failed to load file {}".format(filepath))
         return
-    plotRSS(data, cminmax=(-85, -30))
+
+    print('--- plotting for {}...'.format(filepath))
+
+    plotSingleImage(
+        data,
+        fp=filepath,
+        display=args.showimg,
+        dump=args.dump,
+        cminmax=(-80, -40),
+        title=os.path.basename(filepath).replace('.pickle', '')
+    )
 
 
 def main(args):
@@ -40,7 +50,7 @@ def main(args):
 
     # loop through and plot
     for i in range(len(filepaths)):
-        print("- plotting file {}".format(filepaths[i]))
+        print("- processing file {}".format(filepaths[i]))
         plotSingle(filepaths[i], args)
 
 
@@ -58,6 +68,20 @@ if __name__ == '__main__':
         default=None,
         help='input folderpath for many pickles'
     )
+    p.add_argument(
+        '--dump',
+        dest='dump',
+        action='store_true',
+        default=False,
+        help='input folderpath for many pickles'
+    )
+    p.add_argument(
+        '--no-disp',
+        dest='showimg',
+        action='store_false',
+        default=True,
+        help='input folderpath for many pickles'
+    )
     try:
         args = p.parse_args()
     except BaseException as e: