Add initial implementation of tfplot.contrib (#11)

Add three functions:

- tfplot.contrib.probmap
- tfplot.contrib.probmap_simple
- tfplot.contrib.batch

Author: wookayin

docs/api/contrib.rst

@@ -0,0 +1,10 @@
+.. _api_tfplot_summary:
+
+:mod:`tfplot.summary`
+---------------------
+
+.. contents::
+    :local:
+
+.. automodule:: tfplot.contrib
+    :members:

docs/api/index.rst

@@ -8,4 +8,5 @@ API References
 
     tfplot
     figure
+    contrib
     summary

tfplot/__init__.py

@@ -3,7 +3,8 @@
 from __future__ import print_function
 
 import matplotlib
-matplotlib.use('Agg')
+if not matplotlib.rcParams.get('backend', None):
+    matplotlib.use('Agg')
 
 from .ops import plot, plot_many
 from .wrapper import wrap, wrap_axesplot

tfplot/contrib.py

@@ -0,0 +1,81 @@
+'''Some predefined plot functions.'''
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+from .wrapper import autowrap
+
+
+__all__ = (
+    'probmap',
+    'probmap_simple',
+    'batch',
+)
+
+
+@autowrap
+def probmap(x, cmap='jet', colorbar=True,
+            vmin=None, vmax=None, axis=True, ax=None):
+    '''
+    Display a heatmap in color. The resulting op will be a RGBA image Tensor.
+
+    Args:
+      x: A 2-D image-like tensor to draw.
+      cmap: Matplotlib colormap. Defaults 'jet'
+      axis: If True (default), x-axis and y-axis will appear.
+      colorbar: If True (default), a colorbar will be placed on the right.
+      vmin: A scalar. Minimum value of the range. See ``matplotlib.axes.Axes.imshow``.
+      vmax: A scalar. Maximum value of the range. See ``matplotlib.axes.Axes.imshow``.
+
+    Returns:
+        A `uint8` `Tensor` of shape ``(?, ?, 4)`` containing the resulting plot.
+    '''
+    assert ax is not None, "autowrap did not set ax"
+
+    axim = ax.imshow(x, cmap=cmap, vmin=vmin, vmax=vmax)
+    if colorbar:
+        ax.figure.colorbar(axim)
+    if not axis:
+        ax.axis('off')
+
+    if not axis and not colorbar:
+        ax.figure.subplots_adjust(0, 0, 1, 1)
+    else:
+        ax.figure.tight_layout()
+
+
+def probmap_simple(x, **kwargs):
+    '''
+    Display a heatmap in color, but only displays the image content.
+    The resulting op will be a RGBA image Tensor.
+
+    It reduces to ``probmap`` having `colorbar` and `axis` off.
+    See the documentation of ``probmap`` for available arguments.
+    '''
+    # pylint: disable=unexpected-keyword-arg
+    return probmap(x,
+                   colorbar=kwargs.pop('colorbar', False),
+                   axis=kwargs.pop('axis', False),
+                   figsize=kwargs.pop('figsize', (3, 3)),
+                   **kwargs)
+    # pylint: enable=unexpected-keyword-arg
+
+
+def batch(func):
+    '''
+    Make an autowrapped plot function (... -> RGBA tf.Tensor) work in a batch
+    manner.
+
+    Example:
+
+      >>> p
+      Tensor("p:0", shape=(batch_size, 16, 16, 4), dtype=uint8)
+      >>> tfplot.contrib.batch(tfplot.contrib.probmap)(p)
+      Tensor("probmap/PlotImages:0", shape=(batch_size, ?, ?, 4), dtype=uint8)
+    '''
+    if not hasattr(func, '__unwrapped__'):
+        raise ValueError("The given function is not wrapped with tfplot.autowrap()!")
+
+    func = func.__unwrapped__
+    return autowrap(func, batch=True)

tfplot/wrapper.py

@@ -307,11 +307,16 @@ def _wrapped_plot_fn(*args, **kwargs_call):
 
         return ret
 
+    # return the wrapper (a factory of Tensor)
     _wrapped_fn = wrap(_wrapped_plot_fn, batch=batch, name=name)  # TODO kwargs
 
     _wrapped_fn.__name__ = 'autowrap[%s]' % plot_func.__name__
     if hasattr(plot_func, '__qualname__'):
         _wrapped_fn.__qualname__ = 'autowrap[%s.%s]' % (plot_func.__module__, plot_func.__qualname__)
+
+    # expose the unwrapped python function as well
+    _wrapped_fn.__unwrapped__ = plot_func
+
     return _wrapped_fn