dump vars and visualizations

opt-requirements.txt

@@ -5,3 +5,4 @@ h5py
 pyzmq
 tornado; python_version < '3.0'
 lmdb
+matplotlib

scripts/checkpoint-manipulate.py

@@ -6,13 +6,24 @@
 
 import numpy as np
 from tensorpack.tfutils.varmanip import dump_chkpt_vars
+from tensorpack.utils import logger
 import tensorflow as tf
 import sys
+import argparse
 
-model_path = sys.argv[1]
-reader = tf.train.NewCheckpointReader(model_path)
-var_names = reader.get_variable_to_shape_map().keys()
-result = {}
-for n in var_names:
-    result[n] = reader.get_tensor(n)
-import IPython as IP; IP.embed(config=IP.terminal.ipapp.load_default_config())
+parser = argparse.ArgumentParser()
+parser.add_argument('checkpoint')
+parser.add_argument('--dump', help='dump to an npy file')
+parser.add_argument('--shell', action='store_true', help='start a shell with the params')
+args = parser.parse_args()
+
+if args.checkpoint.endswith('.npy'):
+    params = np.load(args.checkpoint).item()
+else:
+    params = dump_chkpt_vars(args.checkpoint)
+logger.info("Variables in the checkpoint:")
+logger.info(str(params.keys()))
+if args.dump:
+    np.save(args.dump, params)
+if args.shell:
+    import IPython as IP; IP.embed(config=IP.terminal.ipapp.load_default_config())

tensorpack/tfutils/varmanip.py

@@ -87,13 +87,11 @@ the same name".format(v.name))
     logger.info(str(result.keys()))
     np.save(path, result)
 
-def dump_chkpt_vars(model_path, output):
-    """ Dump all variables from a checkpoint """
+def dump_chkpt_vars(model_path):
+    """ Dump all variables from a checkpoint to a dict"""
     reader = tf.train.NewCheckpointReader(model_path)
     var_names = reader.get_variable_to_shape_map().keys()
     result = {}
     for n in var_names:
         result[n] = reader.get_tensor(n)
-    logger.info("Variables to save to {}:".format(output))
-    logger.info(str(result.keys()))
-    np.save(output, result)
+    return result

tensorpack/utils/viz.py

@@ -4,6 +4,37 @@
 # Credit: zxytim
 
 import numpy as np
+import io
+import cv2
+
+try:
+    import matplotlib.pyplot as plt
+except ImportError:
+    pass
+
+__all__ = ['pyplot2img', 'build_patch_list', 'pyplot_viz']
+
+def pyplot2img(plt):
+    buf = io.BytesIO()
+    plt.axis('off')
+    plt.savefig(buf, format='png', bbox_inches='tight', pad_inches=0)
+    buf.seek(0)
+    rawbuf = np.frombuffer(buf.getvalue(), dtype='uint8')
+    im = cv2.imdecode(rawbuf, cv2.IMREAD_COLOR)
+    buf.close()
+    return im
+
+def pyplot_viz(img, shape=None):
+    """ use pyplot to visualize the image
+        Note: this is quite slow. and the returned image will have a border
+    """
+    plt.clf()
+    plt.axes([0,0,1,1])
+    plt.imshow(img)
+    ret = pyplot2img(plt)
+    if shape is not None:
+        ret = cv2.resize(ret, shape)
+    return ret
 
 def minnone(x, y):
     if x is None: x = y
@@ -25,8 +56,10 @@ def build_patch_list(patch_list,
         np.random.shuffle(patch_list)
     ph, pw = patch_list.shape[1:3]
     mh, mw = max(max_height, ph + border), max(max_width, pw + border)
-    nr_row = minnone(nr_row, max_height / (ph + border))
-    nr_col = minnone(nr_col, max_width / (pw + border))
+    if nr_row is None:
+        nr_row = minnone(nr_row, max_height / (ph + border))
+    if nr_col is None:
+        nr_col = minnone(nr_col, max_width / (pw + border))
 
     canvas = np.zeros((nr_row * (ph + border) - border,
              nr_col * (pw + border) - border,