device.py

# -*- coding: utf-8 -*-
# Copyright (c) 2017 Richard Hull and contributors
# See LICENSE.rst for details.

import os
import sys
import atexit
import logging
import string
import curses
import collections
try:
    from StringIO import StringIO
except ImportError:
    from io import StringIO

from PIL import Image, ImageFont, ImageDraw

from luma.core.device import device
from luma.core.serial import noop
from luma.emulator.render import transformer
from luma.emulator.clut import rgb2short


logger = logging.getLogger(__name__)

__all__ = ["capture", "gifanim", "pygame", "asciiart"]


class emulator(device):
    """
    Base class for emulated display driver classes
    """
    def __init__(self, width, height, rotate, mode, transform, scale):
        super(emulator, self).__init__(serial_interface=noop())
        try:
            import pygame
        except ImportError:
            raise RuntimeError("Emulator requires pygame to be installed")
        self._pygame = pygame
        self.capabilities(width, height, rotate, mode)
        self.scale = 1 if transform == "none" else scale
        self._transform = getattr(transformer(pygame, width, height, scale),
                                  "none" if scale == 1 else transform)

    def cleanup(self):
        pass

    def to_surface(self, image, alpha=1.0):
        """
        Converts a :py:mod:`PIL.Image` into a :class:`pygame.Surface`,
        transforming it according to the ``transform`` and ``scale``
        constructor arguments.
        """
        assert(0.0 <= alpha <= 1.0)
        if alpha < 1.0:
            im = image.convert("RGBA")
            black = Image.new(im.mode, im.size, "black")
            im = Image.blend(black, im, alpha)
        else:
            im = image.convert("RGB")

        mode = im.mode
        size = im.size
        data = im.tobytes()
        del im

        surface = self._pygame.image.fromstring(data, size, mode)
        return self._transform(surface)


class capture(emulator):
    """
    Pseudo-device that acts like a physical display, except that it writes the
    image to a numbered PNG file when the :func:`display` method is called.
    Supports 24-bit color depth.
    """
    def __init__(self, width=128, height=64, rotate=0, mode="RGB",
                 transform="scale2x", scale=2, file_template="luma_{0:06}.png",
                 **kwargs):
        super(capture, self).__init__(width, height, rotate, mode, transform, scale)
        self._count = 0
        self._file_template = file_template

    def display(self, image):
        """
        Takes a :py:mod:`PIL.Image` and dumps it to a numbered PNG file.
        """
        assert(image.size == self.size)

        self._count += 1
        filename = self._file_template.format(self._count)
        image = self.preprocess(image)
        logger.debug("Writing: {0}".format(filename))
        image.save(filename)


class gifanim(emulator):
    """
    Pseudo-device that acts like a physical display, except that it collects
    the images when the :func:`display` method is called, and on exit,
    assembles them into an animated GIF image. Supports 24-bit color depth,
    albeit with an indexed color palette.
    """
    def __init__(self, width=128, height=64, rotate=0, mode="RGB",
                 transform="scale2x", scale=2, filename="luma_anim.gif",
                 duration=0.01, loop=0, max_frames=None, **kwargs):
        super(gifanim, self).__init__(width, height, rotate, mode, transform, scale)
        self._images = []
        self._count = 0
        self._max_frames = max_frames
        self._filename = filename
        self._loop = loop
        self._duration = int(duration * 1000)
        atexit.register(self.write_animation)

    def display(self, image):
        """
        Takes an image, scales it according to the nominated transform, and
        stores it for later building into an animated GIF.
        """
        assert(image.size == self.size)

        image = self.preprocess(image)
        surface = self.to_surface(image)
        rawbytes = self._pygame.image.tostring(surface, "RGB", False)
        im = Image.frombytes("RGB", (self._w * self.scale, self._h * self.scale), rawbytes)
        self._images.append(im)

        self._count += 1
        logger.debug("Recording frame: {0}".format(self._count))

        if self._max_frames and self._count >= self._max_frames:
            sys.exit(0)

    def write_animation(self):
        if len(self._images) > 0:
            logger.debug("Please wait... building animated GIF")
            with open(self._filename, "w+b") as fp:
                self._images[0].save(fp, save_all=True, loop=self._loop,
                                     duration=self._duration,
                                     append_images=self._images[1:],
                                     optimize=True, format="GIF")

            file_size = os.stat(self._filename).st_size
            logger.debug("Wrote {0} frames to file: {1} ({2} bytes)".format(
                self._count, self._filename, file_size))


class pygame(emulator):
    """
    Pseudo-device that acts like a physical display, except that it renders
    to a displayed window. The frame rate is limited to 60FPS (much faster
    than a Raspberry Pi can achieve, but this can be overridden as necessary).
    Supports 24-bit color depth.

    :mod:`pygame` is used to render the emulated display window, and it's
    event loop is checked to see if the ESC key was pressed or the window
    was dismissed: if so :func:`sys.exit()` is called.
    """
    def __init__(self, width=128, height=64, rotate=0, mode="RGB", transform="scale2x",
                 scale=2, frame_rate=60, **kwargs):
        super(pygame, self).__init__(width, height, rotate, mode, transform, scale)
        self._pygame.init()
        self._pygame.font.init()
        self._pygame.display.set_caption("Luma.core Display Emulator")
        self._clock = self._pygame.time.Clock()
        self._fps = frame_rate
        self._screen = None
        self._last_image = Image.new(mode, (width, height))
        self._contrast = 1.0

    def _abort(self):
        keystate = self._pygame.key.get_pressed()
        return keystate[self._pygame.K_ESCAPE] or self._pygame.event.peek(self._pygame.QUIT)

    def display(self, image):
        """
        Takes a :py:mod:`PIL.Image` and renders it to a pygame display surface.
        """
        assert(image.size == self.size)
        self._last_image = image

        image = self.preprocess(image)
        self._clock.tick(self._fps)
        self._pygame.event.pump()

        if self._abort():
            self._pygame.quit()
            sys.exit()

        surface = self.to_surface(image, alpha=self._contrast)
        if self._screen is None:
            self._screen = self._pygame.display.set_mode(surface.get_size())
        self._screen.blit(surface, (0, 0))
        self._pygame.display.flip()

    def show(self):
        self.contrast(0xFF)

    def hide(self):
        self.contrast(0x00)

    def contrast(self, value):
        assert(0 <= value <= 255)
        self._contrast = value / 255.0
        self.display(self._last_image)


class asciiart(emulator):
    """
    Pseudo-device that acts like a physical display, except that it converts the
    image to display into an ASCII-art representation and downscales colors to
    match the xterm-256 color scheme. Supports 24-bit color depth.

    This device takes hold of the terminal window (using curses), and any output
    for sysout and syserr is captured and stored, and is replayed when the
    cleanup method is called.

    Loosely based on https://github.com/ajalt/pyasciigen/blob/master/asciigen.py

    .. versionadded:: 0.2.0
    """
    def __init__(self, width=128, height=64, rotate=0, mode="RGB", transform="scale2x",
                 scale=2, **kwargs):

        super(asciiart, self).__init__(width, height, rotate, mode, transform, scale)
        self._stdscr = curses.initscr()
        curses.start_color()
        curses.use_default_colors()
        for i in range(0, curses.COLORS):
            curses.init_pair(i, i, -1)
        curses.noecho()
        curses.cbreak()

        # Capture all stdout, stderr
        self._old_stdX = (sys.stdout, sys.stderr)
        self._captured = (StringIO(), StringIO())
        sys.stdout, sys.stderr = self._captured

        # Sort printable characters according to the number of black pixels present.
        # Don't use string.printable, since we don't want any whitespace except spaces.
        charset = (string.letters + string.digits + string.punctuation + "  ")
        self._chars = list(reversed(sorted(charset, key=self._char_density)))
        self._char_width, self._char_height = ImageFont.load_default().getsize("X")
        self._contrast = 1.0
        self._last_image = Image.new(mode, (width, height))

    def _char_density(self, c, font=ImageFont.load_default()):
        """
        Count the number of black pixels in a rendered character.
        """
        image = Image.new('1', font.getsize(c), color=255)
        draw = ImageDraw.Draw(image)
        draw.text((0, 0), c, fill="white", font=font)
        return collections.Counter(image.getdata())[0]   # 0 is black

    def _generate_art(self, image, width, height):
        """
        Return an iterator that produces the ascii art.
        """
        # Characters aren't square, so scale the output by the aspect ratio of a charater
        height = int(height * self._char_width / float(self._char_height))
        image = image.resize((width, height), Image.ANTIALIAS).convert("RGB")

        for (r, g, b) in image.getdata():
            greyscale = int(0.299 * r + 0.587 * g + 0.114 * b)
            ch = self._chars[int(greyscale / 255. * (len(self._chars) - 1) + 0.5)]
            yield (ch, rgb2short(r, g, b))

    def display(self, image):
        """
        Takes a :py:mod:`PIL.Image` and renders it to the current terminal as
        ASCII-art.
        """
        assert(image.size == self.size)
        self._last_image = image

        surface = self.to_surface(self.preprocess(image), alpha=self._contrast)
        rawbytes = self._pygame.image.tostring(surface, "RGB", False)
        image = Image.frombytes("RGB", (self._w * self.scale, self._h * self.scale), rawbytes)

        scr_width = self._stdscr.getmaxyx()[1]
        scale = float(scr_width) / image.width

        self._stdscr.erase()
        self._stdscr.move(0, 0)
        try:
            for (ch, color) in self._generate_art(image, int(image.width * scale), int(image.height * scale)):
                self._stdscr.addch(ch, curses.color_pair(color))

        except curses.error:
            # End of screen reached
            pass

        self._stdscr.refresh()

    def show(self):
        self.contrast(0xFF)

    def hide(self):
        self.contrast(0x00)

    def contrast(self, value):
        assert(0 <= value <= 255)
        self._contrast = value / 255.0
        self.display(self._last_image)

    def cleanup(self):
        super(asciiart, self).cleanup()

        # Stty sane
        curses.nocbreak()
        curses.echo()
        curses.endwin()

        # Restore stdout & stderr, then print out captured
        sys.stdout, sys.stderr = self._old_stdX
        sys.stdout.write(self._captured[0].getvalue())
        sys.stdout.flush()
        sys.stderr.write(self._captured[1].getvalue())
        sys.stderr.flush()