dito.conversion

  1"""
  2This submodule provides functionality for the conversion of NumPy arrays to other formats and vice versa.
  3"""
  4import io
  5
  6import cv2
  7import numpy as np
  8
  9import dito.core
 10import dito.exceptions
 11
 12
 13#
 14# matplotlib
 15#
 16
 17
 18def fig_to_image(fig, size=(800, 600), savefig_kwargs=None):
 19    """
 20    Convert a Matplotlib figure to a NumPy image array.
 21
 22    Parameters
 23    ----------
 24    fig : matplotlib.figure.Figure
 25        The Matplotlib figure to convert.
 26    size : tuple of int, optional
 27        Desired output image size in pixels as (width, height). Default is (800, 600).
 28    savefig_kwargs : dict, optional
 29        Additional keyword arguments passed to `fig.savefig`. Can override default options
 30        like facecolor or transparency.
 31
 32    Returns
 33    -------
 34    numpy.ndarray
 35        Image array in HWC format (height, width, 3), with BGR channels.
 36
 37    Notes
 38    -----
 39    Removes alpha channel from the saved PNG to avoid transparency artifacts,
 40    following known Matplotlib behavior (see issue #14339).
 41
 42    Examples
 43    --------
 44    >>> (fig, ax) = plt.subplots()                  # doctest: +SKIP
 45    >>> ax.plot([0, 1], [0, 1])                     # doctest: +SKIP
 46    >>> image = fig_to_image(fig, size=(400, 300))  # doctest: +SKIP
 47    >>> image.shape                                 # doctest: +SKIP
 48    (300, 400, 3)                                   # doctest: +SKIP
 49    """
 50
 51    # set figure size in pixels
 52    (width, height) = size
 53    dpi = width / fig.get_size_inches()[0]
 54    fig.set_size_inches(width / dpi, height / dpi)
 55
 56    # save figure to buffer
 57    savefig_kwargs_merged = dict(
 58        facecolor="white",
 59        transparent=False,
 60    )
 61    if savefig_kwargs is not None:
 62        savefig_kwargs_merged.update(savefig_kwargs)
 63    buffer = io.BytesIO()
 64    fig.savefig(buffer, format="png", dpi=dpi, **savefig_kwargs_merged)
 65
 66    # read image from buffer
 67    buffer.seek(0)
 68    png_bytes = np.frombuffer(buffer.getvalue(), dtype=np.uint8)
 69    image = cv2.imdecode(buf=png_bytes, flags=cv2.IMREAD_UNCHANGED)
 70
 71    # remove alpha channel (see https://github.com/matplotlib/matplotlib/issues/14339)
 72    if (len(image.shape) == 3) and (image.shape[2] == 4):
 73        image = image[:, :, :3]
 74
 75    return image
 76
 77
 78#
 79# PySide6
 80#
 81
 82
 83def to_PySide6_QPixmap_format(image):
 84    """
 85    Determine the QImage.Format which is compatible with the given image.
 86
 87    Parameters
 88    ----------
 89    image : np.ndarray
 90        The input image.
 91
 92    Returns
 93    -------
 94    PySide6.QtGui.QImage.Format
 95        The QImage.Format which is compatible with the given image.
 96
 97    Raises
 98    ------
 99    ImportError
100        If PySide6 is not installed.
101    dito.exceptions.ConversionError
102        If the given image cannot be converted to a compatible QImage.Format.
103    """
104    import PySide6.QtGui
105
106    dtype = image.dtype
107    if dito.core.is_gray(image):
108        if dtype == np.uint8:
109            return PySide6.QtGui.QImage.Format_Grayscale8
110        elif dtype == np.uint16:
111            return PySide6.QtGui.QImage.Format_Grayscale16
112        else:
113            raise dito.exceptions.ConversionError("Conversion of grayscale image with dtype '{}' to QPixmap is not supported".format(dtype))
114
115    elif dito.core.is_color(image):
116        if dtype == np.uint8:
117            return PySide6.QtGui.QImage.Format_BGR888
118        else:
119            raise dito.exceptions.ConversionError("Conversion of color image with dtype '{}' to QPixmap is not supported".format(dtype))
120
121    else:
122        raise dito.exceptions.ConversionError("Conversion image with shape {} to QPixmap is not supported".format(image.shape))
123
124
125def to_PySide6_QImage(image):
126    """
127    Convert a numpy.ndimage to PySide6.QtGui.QImage.QImage.
128
129    Parameters
130    ----------
131    image : np.ndarray
132        The input image.
133
134    Returns
135    -------
136    PySide6.QtGui.QImage
137        The QImage representation of the input image.
138
139    Raises
140    ------
141    ImportError
142        If PySide6 is not installed.
143    """
144    import PySide6.QtGui
145    return PySide6.QtGui.QImage(
146        np.require(image, requirements="C"),
147        image.shape[1],
148        image.shape[0],
149        to_PySide6_QPixmap_format(image),
150    )
151
152
153def to_PySide6_QPixmap(image):
154    """
155    Convert a numpy.ndimage to PySide6.QtGui.QImage.QPixmap.
156
157    Parameters
158    ----------
159    image : np.ndarray
160        The input image.
161
162    Returns
163    -------
164    PySide6.QtGui.QPixmap
165        The QPixmap representation of the input image.
166
167    Raises
168    ------
169    ImportError
170        If PySide6 is not installed.
171    """
172    import PySide6.QtGui
173    q_image = to_PySide6_QImage(image)
174    return PySide6.QtGui.QPixmap(q_image)