Merge pull request #2791 from ericpre/output_size_save_image

Output size save image

doc/user_guide/io.rst

@@ -588,25 +588,70 @@ HyperSpy can read and write data to `all the image formats
 <https://imageio.readthedocs.io/en/stable/formats.html>`_ supported by
 `imageio`, which uses the Python Image Library  (PIL/pillow).
 This includes png, pdf, gif, etc.
+It is important to note that these image formats only support 8-bit files, and
+therefore have an insufficient dynamic range for most scientific applications.
+It is therefore highly discouraged to use any general image format (with the
+exception of :ref:`tiff-format` which uses another library) to store data for
+analysis purposes.
+
+Extra saving arguments
+^^^^^^^^^^^^^^^^^^^^^^
+
+- ``scalebar`` (bool, optional): Export the image with a scalebar. Default
+  is False.
+- ``scalebar_kwds`` (dict, optional): dictionary of keyword arguments for the
+  scalebar. Useful to set formattiong, location, etc. of the scalebar. See the
+  `matplotlib-scalebar <https://pypi.org/project/matplotlib-scalebar/>`_
+  documentation for more information.
+- ``output_size`` : (int, tuple of length 2 or None, optional): the output size 
+  of the image in pixels:
+
+  * if ``int``, defines the width of the image, the height is
+    determined from the aspect ratio of the image.
+  * if ``tuple`` of length 2, defines the width and height of the
+    image. Padding with white pixels is used to maintain the aspect
+    ratio of the image.
+  * if ``None``, the size of the data is used.
+
+  For output sizes larger than the data size, "nearest" interpolation is
+  used by default and this behaviour can be changed through the
+  ``imshow_kwds`` dictionary.
+
+- ``imshow_kwds`` (dict, optional):  Keyword arguments dictionary for
+  :py:func:`~.matplotlib.pyplot.imshow`.
+- ``**kwds`` : keyword arguments supported by the individual file
+  writers as documented at
+  https://imageio.readthedocs.io/en/stable/formats.html when exporting
+  an image without scalebar. When exporting with a scalebar, the keyword
+  arguments are passed to the `pil_kwargs` dictionary of
+  :py:func:`matplotlib.pyplot.savefig`
+
 
 When saving an image, a scalebar can be added to the image and the formatting,
-location, etc. of the scalebar can be set using the ``scalebar_kwds`` arguments
-- see the `matplotlib-scalebar <https://pypi.org/project/matplotlib-scalebar/>`_
-documentation for more information.
+location, etc. of the scalebar can be set using the ``scalebar_kwds``
+arguments:
 
 .. code-block:: python
 
     >>> s.save('file.jpg', scalebar=True)
     >>> s.save('file.jpg', scalebar=True, scalebar_kwds={'location':'lower right'})
 
-When saving an image, keyword arguments can be passed to the corresponding
-pillow file writer.
+In the example above, the image is created using
+:py:func:`~.matplotlib.pyplot.imshow`, and additional keyword arguments can be
+passed to this function using ``imshow_kwds``. For example, this can be used
+to save an image displayed using a matplotlib colormap:
+
+.. code-block:: python
+
+    >>> s.save('file.jpg', imshow_kwds=dict(cmap='viridis'))
+
+
+The resolution of the exported image can be adjusted:
+
+.. code-block:: python
+
+    >>> s.save('file.jpg', output_size=512)
 
-It is important to note that these image formats only support 8-bit files, and
-therefore have an insufficient dynamic range for most scientific applications.
-It is therefore highly discouraged to use any general image format (with the
-exception of :ref:`tiff-format` which uses another library) to store data for
-analysis purposes.
 
 .. _tiff-format:
 
@@ -822,10 +867,10 @@ Extra saving arguments
 - ``intensity_scaling`` : in case the dataset that needs to be saved does not
   have the `np.uint8` data type, casting to this datatype without intensity
   rescaling results in overflow errors (default behavior). This option allows
-  you to perform linear intensity scaling of the images prior to saving the 
+  you to perform linear intensity scaling of the images prior to saving the
   data. The options are:
-  
-  - `'dtype'`: the limits of the datatype of the dataset, e.g. 0-65535 for 
+
+  - `'dtype'`: the limits of the datatype of the dataset, e.g. 0-65535 for
     `np.uint16`, are mapped onto 0-255 respectively. Does not work for `float`
     data types.
   - `'minmax'`: the minimum and maximum in the dataset are mapped to 0-255.
@@ -835,7 +880,7 @@ Extra saving arguments
 - ``navigator_signal``: the BLO file also stores a virtual bright field (VBF) image which
   behaves like a navigation signal in the ASTAR software. By default this is
   set to `'navigator'`, which results in the default :py:attr:`navigator` signal to be used.
-  If this signal was not calculated before (e.g. by calling :py:meth:`~.signal.BaseSignal.plot`), it is 
+  If this signal was not calculated before (e.g. by calling :py:meth:`~.signal.BaseSignal.plot`), it is
   calculated when :py:meth:`~.signal.BaseSignal.save` is called, which can be time consuming.
   Alternatively, setting the argument to `None` will result in a correctly sized
   zero array to be used. Finally, a custom ``Signal2D`` object can be passed,

hyperspy/datasets/artificial_data.py

@@ -41,11 +41,8 @@ def get_low_loss_eels_signal(add_noise=True, random_state=None):
     Parameters
     ----------
     %s
-    %s
 
-    Returns
-    -------
-    artificial_low_loss_signal : :py:class:`~hyperspy._signals.eels.EELSSpectrum`
+    %s
 
     Example
     -------
@@ -84,7 +81,8 @@ def get_low_loss_eels_signal(add_noise=True, random_state=None):
                                      RETURNS_DOCSTRING)
 
 
-def get_core_loss_eels_signal(add_powerlaw=False, add_noise=True, random_state=None):
+def get_core_loss_eels_signal(add_powerlaw=False, add_noise=True,
+                              random_state=None):
     """Get an artificial core loss electron energy loss spectrum.
 
     Similar to a Mn-L32 edge from a perovskite oxide.
@@ -209,7 +207,8 @@ def get_low_loss_eels_line_scan_signal(add_noise=True, random_state=None):
                                                RETURNS_DOCSTRING)
 
 
-def get_core_loss_eels_line_scan_signal(add_powerlaw=False, add_noise=True, random_state=None):
+def get_core_loss_eels_line_scan_signal(add_powerlaw=False, add_noise=True,
+                                        random_state=None):
     """Get an artificial core loss electron energy loss line scan spectrum.
 
     Similar to a Mn-L32 and Fe-L32 edge from a perovskite oxide.
@@ -279,7 +278,8 @@ def get_core_loss_eels_line_scan_signal(add_powerlaw=False, add_noise=True, rand
                                                 RETURNS_DOCSTRING)
 
 
-def get_core_loss_eels_model(add_powerlaw=False, add_noise=True, random_state=None):
+def get_core_loss_eels_model(add_powerlaw=False, add_noise=True,
+                             random_state=None):
     """Get an artificial core loss electron energy loss model.
 
     Similar to a Mn-L32 edge from a perovskite oxide.
@@ -364,12 +364,10 @@ def get_luminescence_signal(navigation_dimension=0,
     uniform: bool.
         return uniform (wavelength) or non-uniform (energy) spectrum
     add_baseline : bool
-            If true, adds a constant baseline to the spectrum. Conversion to
-            energy representation will turn the constant baseline into inverse
-            powerlaw.
+        If true, adds a constant baseline to the spectrum. Conversion to
+        energy representation will turn the constant baseline into inverse
+        powerlaw.
     %s
-    random_state: None or int
-        initialise state of the random number generator
 
     Example
     -------

hyperspy/io_plugins/image.py

@@ -21,8 +21,8 @@
 
 from imageio import imread, imwrite
 from matplotlib.figure import Figure
-import traits.api as t
 import pint
+import traits.api as t
 
 from hyperspy.misc import rgb_tools
 
@@ -35,78 +35,138 @@
                    'msp', 'pcx', 'ppm', "pbm", "pgm", 'xbm', 'spi', ]
 default_extension = 0  # png
 # Writing capabilities
-writes = [(2, 0), ]
+writes = [(2, 0), (0, 2)]
 non_uniform_axis = False
 # ----------------------
 
 _ureg = pint.UnitRegistry()
 _logger = logging.getLogger(__name__)
 
 
-def file_writer(filename, signal, scalebar=False,
-                scalebar_kwds={'box_alpha':0.75, 'location':'lower left'},
-                **kwds):
-    """Writes data to any format supported by PIL
+def file_writer(filename, signal, scalebar=False, scalebar_kwds=None,
+                output_size=None, imshow_kwds=None, **kwds):
+    """Writes data to any format supported by pillow. When ``output_size``
+    or ``scalebar`` or ``imshow_kwds`` is used,
+    :py:func:`~.matplotlib.pyplot.imshow` is used to generate a figure.
 
     Parameters
     ----------
     filename: {str, pathlib.Path, bytes, file}
         The resource to write the image to, e.g. a filename, pathlib.Path or
-        file object, see the docs for more info. The file format is defined by 
+        file object, see the docs for more info. The file format is defined by
         the file extension that is any one supported by imageio.
     signal: a Signal instance
     scalebar : bool, optional
-        Export the image with a scalebar.
-    scalebar_kwds : dict
+        Export the image with a scalebar. Default is False.
+    scalebar_kwds : dict, optional
         Dictionary of keyword arguments for the scalebar. Useful to set
         formattiong, location, etc. of the scalebar. See the documentation of
         the 'matplotlib-scalebar' library for more information.
-    **kwds: keyword arguments
+    output_size : {tuple of length 2, int, None}, optional
+        The output size of the image in pixels (width, height):
+
+            * if *int*, defines the width of the image, the height is
+              determined from the aspec ratio of the image
+            * if *tuple of length 2*, defines the width and height of the
+              image. Padding with white pixels is used to maintain the aspect
+              ratio of the image.
+            * if *None*, the size of the data is used.
+
+        For output size larger than the data size, "nearest" interpolation is
+        used by default and this behaviour can be changed through the
+        *imshow_kwds* dictionary. Default is None.
+
+    imshow_kwds : dict, optional
+        Keyword arguments dictionary for :py:func:`~.matplotlib.pyplot.imshow`.
+    **kwds : keyword arguments, optional
         Allows to pass keyword arguments supported by the individual file
-        writers as documented at https://imageio.readthedocs.io/en/stable/formats.html
+        writers as documented at
+        https://imageio.readthedocs.io/en/stable/formats.html when exporting
+        an image without scalebar. When exporting with a scalebar, the keyword
+        arguments are passed to the `pil_kwargs` dictionary of
+        :py:func:`~matplotlib.pyplot.savefig`
 
     """
     data = signal.data
+
+    if scalebar_kwds is None:
+        scalebar_kwds = dict()
+    scalebar_kwds.setdefault('box_alpha', 0.75)
+    scalebar_kwds.setdefault('location', 'lower left')
+
     if rgb_tools.is_rgbx(data):
         data = rgb_tools.rgbx2regular_array(data)
+
     if scalebar:
         try:
             from matplotlib_scalebar.scalebar import ScaleBar
-            export_scalebar = True
         except ImportError:  # pragma: no cover
-            export_scalebar = False
+            scalebar = False
             _logger.warning("Exporting image with scalebar requires the "
                             "matplotlib-scalebar library.")
+
+    if scalebar or output_size or imshow_kwds:
         dpi = 100
-        fig = Figure(figsize=[v/dpi for v in signal.axes_manager.signal_shape],
-                     dpi=dpi)
 
-        try:
-            # List of format supported by matplotlib
-            supported_format = sorted(fig.canvas.get_supported_filetypes())
-            if os.path.splitext(filename)[1].replace('.', '') not in supported_format:
-                export_scalebar = False
-                _logger.warning("Exporting image with scalebar is supported "
-                                f"only with {', '.join(supported_format)}.")
-        except AttributeError:  # pragma: no cover
-            export_scalebar = False
-            _logger.warning("Exporting image with scalebar requires the "
-                            "matplotlib 3.1 or newer.")
+        if imshow_kwds is None:
+            imshow_kwds = dict()
+        imshow_kwds.setdefault('cmap', 'gray')
+
+        if len(signal.axes_manager.signal_axes) == 2:
+            axes = signal.axes_manager.signal_axes
+        elif len(signal.axes_manager.navigation_axes) == 2:
+            # Use navigation axes
+            axes = signal.axes_manager.navigation_axes
+
+        aspect_ratio = imshow_kwds.get('aspect', None)
+        if not isinstance(aspect_ratio, (int, float)):
+            aspect_ratio = data.shape[0] / data.shape[1]
+
+        if output_size is None:
+            # fall back to image size taking into account aspect_ratio
+            ratio = (1,  aspect_ratio)
+            output_size = [axis.size * r for axis, r in zip(axes, ratio)]
+        elif isinstance(output_size, (int, float)):
+            output_size = [output_size, output_size * aspect_ratio]
+
+        fig = Figure(figsize=[size / dpi for size in output_size], dpi=dpi)
+
+        # List of format supported by matplotlib
+        supported_format = sorted(fig.canvas.get_supported_filetypes())
+        if os.path.splitext(filename)[1].replace('.', '') not in supported_format:
+            if scalebar:
+                raise ValueError("Exporting image with scalebar is supported "
+                                 f"only with {', '.join(supported_format)}.")
+            if output_size:
+                raise ValueError("Setting the output size is only supported "
+                                 f"with {', '.join(supported_format)}.")
 
-    if scalebar and export_scalebar:
+    if scalebar:
+        # Sanity check of the axes
+        # This plugin doesn't support non-uniform axes, we don't need to check
+        # if the axes have a scale attribute
+        if axes[0].scale != axes[1].scale or axes[0].units != axes[1].units:
+            raise ValueError("Scale and units must be the same for each axes "
+                             "to export images with a scale bar.")
+
+    if scalebar or output_size:
         ax = fig.add_axes([0, 0, 1, 1])
         ax.axis('off')
-        ax.imshow(data, cmap='gray')
-
-        # Add scalebar
-        axis = signal.axes_manager.signal_axes[0]
-        if axis.units == t.Undefined:
-            axis.units = "px"
-            scalebar_kwds['dimension'] = "pixel-length"
-        if _ureg.Quantity(axis.units).check('1/[length]'):
-            scalebar_kwds['dimension'] = "si-length-reciprocal"
-        scalebar = ScaleBar(axis.scale, axis.units, **scalebar_kwds)
-        ax.add_artist(scalebar)
+        ax.imshow(data, **imshow_kwds)
+
+        if scalebar:
+            # Add scalebar
+            axis = axes[0]
+            units = axis.units
+            if units == t.Undefined:
+                units = "px"
+                scalebar_kwds['dimension'] = "pixel-length"
+            if _ureg.Quantity(units).check('1/[length]'):
+                scalebar_kwds['dimension'] = "si-length-reciprocal"
+
+            scalebar = ScaleBar(axis.scale, units, **scalebar_kwds)
+            ax.add_artist(scalebar)
+
         fig.savefig(filename, dpi=dpi, pil_kwargs=kwds)
     else:
         imwrite(filename, data, **kwds)