Module eink.image

class EinkGraphics:
    """Provides static methods for reducing images to e-ink display palettes.
    """
    @staticmethod
    def _has_alpha(image):
        """Return whether the specified ``Image`` has an alpha channel."""
        return image.mode[-1] in ['A', 'a']

    @staticmethod
    def _assert_doesnt_have_alpha(image):
        """Raise a ``ValueError`` if the given ``Image`` has an alpha channel.
        """
        if EinkGraphics._has_alpha(image):
            raise ValueError('Alpha channels are not supported')

    @staticmethod
    def round(image, palette=Palette.THREE_BIT_GRAYSCALE):
        """Return the result of rounding the given image to the given palette.

        Return the result of rounding each pixel in the specified
        ``Image`` to roughly the nearest color in the specified
        ``Palette``. This is a format suitable for display on an e-ink
        device. The image may not have an alpha channel.

        This does not perform dithering. See also ``dither``.
        """
        EinkGraphics._assert_doesnt_have_alpha(image)
        if palette._is_grayscale:
            # First convert to grayscale, in order to apply the luminosity
            # transform function L = 0.299 * R + 0.587 * G + 0.114 * B
            grayscale_image = image.convert('L')

            # At time of writing, the "quantize" method rounds each component
            # down to the nearest multiple of four before rounding a pixel to
            # the nearest palette color. We can round pixels more accurately by
            # calling "point" instead.
            return grayscale_image.point(palette._round_lookup_table())
        else:
            return image.convert('RGB').quantize(
                dither=Image.Dither.NONE, palette=palette._image())

    @staticmethod
    def dither(image, palette=Palette.THREE_BIT_GRAYSCALE):
        """Return the result of dithering the given image to the given palette.

        Return the result of converting each pixel in the specified
        ``Image`` to a color in the specified ``Palette`` using
        dithering. This is a format suitable for display on an e-ink
        device. The image may not have an alpha channel.

        Instead of simply rounding each pixel to the nearest color in
        the palette, we use dithering, which represents a region that is
        in between two palette colors by coloring its pixels using a
        combination of those two colors. See
        https://en.wikipedia.org/wiki/Dither .

        Dithering gives the result a speckled appearance that tends to
        look closer to the original image, especially from a distance.
        However, in some cases, a more flatly shaded look might be
        preferable. For example, dithering might be undesirable for
        icons that have large areas of solid shading.
        """
        EinkGraphics._assert_doesnt_have_alpha(image)
        if palette._is_grayscale:
            # First convert to grayscale, in order to apply the luminosity
            # transform function L = 0.299 * R + 0.587 * G + 0.114 * B
            image = image.convert('L')
        return image.convert('RGB').quantize(
            dither=Image.Dither.FLOYDSTEINBERG, palette=palette._image())

class Palette:
    """A color palette for an e-ink device.

    The following palettes are supported:

    * ``Palette.THREE_BIT_GRAYSCALE``: 3-bit grayscale, i.e. eight
      shades of gray.
    * ``Palette.FOUR_BIT_GRAYSCALE``: 4-bit grayscale, i.e. 16
      shades of gray.
    * ``Palette.MONOCHROME``: Black and white.
    * ``Palette.BLACK_WHITE_AND_RED``: Black, white, and red.
    * ``Palette.SEVEN_COLOR``: The palette for color Inkplate devices.
      This has the following seven colors: black, white, red, green,
      blue, yellow, and orange.
    """

    # Private attributes:
    #
    # list<tuple<int, int, int>> _colors - The colors in the palette. Each
    #     color is represented as a tuple of the red, green, and blue
    #     components, in the range [0, 255].
    # Image _image_cache - The cached return value of _image().
    # bool _is_grayscale - Whether the palette consists exclusively of
    #     grayscale colors.
    # string _name - A string identifying the palette in client code. This
    #     consists exclusively of underscores and uppercase letters.
    # list<int> _round_lookup_table_cache - The cached return value of
    #     _round_lookup_table().

    def __init__(self, colors, name):
        """Private initializer."""
        if len(colors) > 256:
            raise ValueError('Palette only supports up to 256 colors')

        self._colors = colors
        self._name = name
        self._round_lookup_table_cache = None
        self._image_cache = None

        self._is_grayscale = True
        for color in colors:
            if color[0] != color[1] or color[0] != color[2]:
                self._is_grayscale = False
                break

    def _round_lookup_table(self):
        """Return a lookup table for rounding to the nearest grayscale color.

        Assume that ``_is_grayscale`` is true. The return value is an
        array of 256 integers. The (i + 1)th element is the red, green,
        and blue component of the palette color nearest to the color
        ``(i, i, i)`` (with ties broken arbitrarily).
        """
        if self._round_lookup_table_cache is None:
            self._round_lookup_table_cache = []
            sorted_colors = sorted([color[0] for color in self._colors])
            index = 0
            for color in range(256):
                if (index + 1 < len(sorted_colors) and
                        sorted_colors[index + 1] - color <
                        color - sorted_colors[index]):
                    index += 1
                self._round_lookup_table_cache.append(sorted_colors[index])
        return self._round_lookup_table_cache

    def _image(self):
        """Return an ``Image`` whose palette is the colors of this palette.

        Return an ``Image`` of mode ``'P'`` whose palette consists of
        the colors of this palette.
        """
        if self._image_cache is None:
            palette = []
            for color in self._colors:
                palette.extend(color)
            self._image_cache = Image.new('P', (1, 1))
            self._image_cache.putpalette(palette)
        return self._image_cache