PIT: export to matplotlib plot with custom right-click menu option.

arpys/pit/imageplot.py

@@ -2,12 +2,15 @@
 
 import logging
 
+import matplotlib.pyplot as plt
 import pyqtgraph as pg
-from numpy import clip, inf, linspace, ndarray
+from matplotlib.colors import ListedColormap
+from numpy import array, clip, inf, linspace, ndarray
 from pyqtgraph import Qt as qt #import QtCore
 from pyqtgraph.graphicsItems.ImageItem import ImageItem
 from pyqtgraph.widgets import PlotWidget, GraphicsView
 
+from arpys.pit.pitviewbox import PITViewBox
 from arpys.pit.utilities import TracedVariable, indexof
 
 logger = logging.getLogger('pit.'+__name__)
@@ -26,6 +29,9 @@ class ImagePlot(pg.PlotWidget) :
     sig_axes_changed   emitted when the axes are updated
     =================  =========================================================
     """
+    # np.array, raw image data
+    image_data = None
+    # pg.ImageItem of *image_data*
     image_item = None
     image_kwargs = {}
     xlim = None
@@ -45,7 +51,8 @@ def __init__(self, image=None, parent=None, background='default',
         name        str; allows giving a name for debug purposes
         ==========  ============================================================
         """
-        super().__init__(parent=parent, background=background, **kwargs) 
+        super().__init__(parent=parent, background=background, 
+                         viewBox=PITViewBox(imageplot=self), **kwargs) 
         self.name = name
 
         # Show top and tight axes by default, but without ticklabels
@@ -81,6 +88,8 @@ def set_image(self, image, *args, **kwargs) :
         kwargs  `ImageItem <pyqtgraph.graphicsItems.ImageItem.ImageItem>`
         ======  ================================================================
         """
+        self.image_data = image
+
         # Convert array to ImageItem
         if isinstance(image, ndarray) :
             image = ImageItem(image, *args, **kwargs)
@@ -212,6 +221,41 @@ def release_viewrange(self) :
                        maxXRange=inf,
                        maxYRange=inf)
 
+    def mpl_export(self) :
+        """ Export the content of this plot to a png image using matplotlib. 
+        The resulting image will have a white background and black ticklabes 
+        and should therefore be more readable than pyqtgraph's native plot
+        export options.
+        """
+        logger.debug('<ImagePlot.mpl_export()>')
+
+        # Get a filename first
+        fd = qt.QtGui.QFileDialog()
+        filename = fd.getSaveFileName()[0]
+        logger.debug('Outfilename: {}'.format(filename))
+
+        # Access data and all relevant display options
+        data = self.image_data.T
+        lut = self.image_item.lut
+
+        # Convert the lookuptable (lut) to a matplotlib colormap
+        lut = lut/lut.max()
+        cmap_array = array([[a[0], a[1], a[2], 1.] for a in lut])
+        cmap = ListedColormap(cmap_array)
+
+        # Temporarily turn matplotlib's interactive mode off
+        was_interactive = plt.isinteractive() 
+        if was_interactive:
+            plt.ioff()
+
+        # Create a matplotlib figure and save it
+        fig, ax = plt.subplots(1)
+        mesh = ax.pcolormesh(self.xscale, self.yscale, data, cmap=cmap)
+        fig.savefig(filename, dpi=300)
+
+        if was_interactive:
+            plt.ion()
+
 class Crosshair() :
     """ Crosshair made up of two InfiniteLines. """
 

arpys/pit/mainwindow.py

@@ -540,9 +540,6 @@ def update_cut(self) :
             return
 
         self.data_handler.cut_data = cut
-        # Convert np.array *cut* to an ImageItem and set it as *cut_plot*'s 
-        # Image
-        cut_image = pg.ImageItem(cut)
         self.cut_plot.set_image(cut, lut=self.lut)
 
     def update_image_data(self) :
@@ -674,7 +671,8 @@ def rotate(self, alpha=0) :
         # Remember that the order in which transformations are applied is 
         # reverted to how they added in the code, i.e. last transform added 
         # in the code will come first (this is the reason we have to 
-        # completely rebuild the transformation instead of just adding a rotation...)
+        # completely rebuild the transformation instead of just adding a 
+        # rotation...)
         transform.reset()
         transform.translate(dx/sx, dy/sy)
         transform.translate(wx/2, wy/2)

arpys/pit/pitviewbox.py

@@ -0,0 +1,33 @@
+"""
+Subclass of ViewBox with custom menu items.
+"""
+import logging
+
+from pyqtgraph.Qt import QtGui
+from pyqtgraph.graphicsItems.ViewBox import ViewBox, ViewBoxMenu
+
+logger = logging.getLogger('pit.'+__name__)
+
+class PITViewBoxMenu(ViewBoxMenu.ViewBoxMenu) :
+    """ Subclass of ViewBoxMenu with custom menu items. """
+    def __init__(self, *args, **kwargs) :
+        super().__init__(*args, **kwargs)
+
+        # Define own menu entries
+        self.mpl_export = QtGui.QAction('MPL export', self)
+
+class PITViewBox(ViewBox) :
+    """
+    Subclass of ViewBox with custom menu items, as defined in `PITViewBoxMenu 
+    <arpys.pit.pitviewbox.PITViewBoxMenu>`.
+    """
+    def __init__(self, imageplot, *args, **kwargs) :
+        super().__init__(*args, **kwargs)
+        logger.debug(imageplot)
+        self.imageplot = imageplot
+
+        # Connect signal handling and make the menu entry appear
+        self.menu = PITViewBoxMenu(self)
+        self.menu.mpl_export.triggered.connect(self.imageplot.mpl_export)
+        self.menu.addAction(self.menu.mpl_export)
+