From 3259127324c560a6d9f1d27779b7ced1ca3571a4 Mon Sep 17 00:00:00 2001
From: HappyZ <zhuyanzi@gmail.com>
Date: Fri, 31 May 2019 14:04:21 -0500
Subject: [PATCH] change functions to plot image

---
 libs/plotting.py | 30 ++++++++++++++++++++++++++++++
 visualize.py     | 30 +++++++++++++++++++++++++++---
 2 files changed, 57 insertions(+), 3 deletions(-)

diff --git a/libs/plotting.py b/libs/plotting.py
index ec3e0f9..479e53b 100644
--- a/libs/plotting.py
+++ b/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()
+
 
diff --git a/visualize.py b/visualize.py
index c4ead29..167825c 100644
--- a/visualize.py
+++ b/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: