propagation_gan/visualize.py

102 lines
2.4 KiB
Python

#!/usr/bin/python
import os
import sys
import pickle
import argparse
from libs.plotting import plotSingleImage
def loadData(filepath):
try:
return pickle.load(open(filepath, "rb"))
except BaseException as e:
print("err: {}".format(e))
pass
return None
def plotSingle(filepath, args):
'''
'''
# load data from file
data = loadData(filepath)
if data is None:
print("err: failed to load file {}".format(filepath))
return
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):
# finding files
filepaths = []
if args.filepath:
filepaths.append(args.filepath)
if args.folderpath:
files = os.listdir(args.folderpath)
filepaths.extend(["{}/{}".format(args.folderpath, file) for file in files if '.pickle' in file])
# loop through and plot
for i in range(len(filepaths)):
print("- processing file {}".format(filepaths[i]))
plotSingle(filepaths[i], args)
if __name__ == '__main__':
p = argparse.ArgumentParser(description='Visualization')
p.add_argument(
'--filepath', '-f',
dest='filepath',
default=None,
help='input filepath for a pickle'
)
p.add_argument(
'--folderpath', '-fd',
dest='folderpath',
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:
print(e)
sys.exit()
if args.filepath is None and args.folderpath is None:
print("at least specify file `-f` or folder `-fd`")
sys.exit()
elif args.filepath is None and not os.path.isdir(args.folderpath):
print("folder {} not exit".format(args.folderpath))
sys.exit()
elif args.folderpath is None and not os.path.isfile(args.filepath):
print("file {} not exit".format(args.filepath))
sys.exit()
main(args)