particularly simple digitizer with matplotlib

digitizer.py

#! /usr/bin/env python
#
# This is a particularly simple digitizer implemented with matplotlib. Intended
# for converting graphs and figures into pixel measurements that can be
# transformed into useful data.
#
# Requires Python 3.4+ because I was playing with asyncio coroutines.
#
# Nat Wilson
# March, 2015

import argparse
import asyncio
import readline
import sys
import matplotlib.pyplot as plt
import numpy as np
from PIL import Image

class Cursor(object):

    def __init__(self, ax):
        self.ax = ax
        self.lx = ax.axhline(color='k', lw=0.5)
        self.ly = ax.axvline(color='k', lw=0.5)
        return

    def mouse_move(self, event):
        if event.inaxes:
            x, y = event.xdata, event.ydata
            self.lx.set_ydata(y)
            self.ly.set_xdata(x)
        return

class State(object):

    def __init__(self):
        self.label = "d"
        self.x = []
        self.y = []
        self.labels = []
        return

    def write(self, fnm):
        with open(fnm, "w") as f:
            for lbl, x, y in zip(self.labels, self.x, self.y):
                f.write("{0} {1} {2}\n".format(lbl, x, y))

class Digitizer(object):

    def __init__(self, state, image, **kw):
        self._state = state
        self._fig = plt.figure()
        self._ax = self._fig.add_axes([0, 0, 1, 1], frameon=False)
        self._ax.set_xticks([])
        self._ax.set_yticks([])
        self._ax.imshow(image, **kw)
        self._lines = []
        self.cid = self._fig.canvas.mpl_connect('button_press_event',
                                                self._clickhandler)
        self.cursor = Cursor(self._ax)
        self._fig.canvas.mpl_connect("motion_notify_event", self.cursor.mouse_move)
        return

    def _clickhandler(self, event):
        # only does anything if zoom/pan not activated
        if self._fig.canvas.manager.toolbar._active is None:
            if event.button == 1:
                self._state.x.append(event.xdata)
                self._state.y.append(event.ydata)
                self._state.labels.append(self._state.label)
                self._lines.append(self._ax.plot(event.xdata, event.ydata, 
                                                   color="r", marker="."))
            elif event.button == 3:
                try:
                    self._state.x.pop()
                    self._state.y.pop()
                    self._state.labels.pop()
                    self._lines[-1][0].remove()
                    self._lines.pop()
                except IndexError:
                    pass
        return

@asyncio.coroutine
def get_cli_cmd(state):
    try:
        ret = input("{0} >> ".format(state.label))
    except EOFError:
        raise KeyboardInterrupt()
    return ret

@asyncio.coroutine
def handle_cli_cmd(state):
    while True:
        cmd = yield from get_cli_cmd(state)
        if cmd.startswith("label "):
            state.label = cmd.split()[1]
        elif cmd.startswith("write "):
            state.write(cmd.split()[1])
            print("{0} written".format(cmd.split()[1]))
        elif cmd.strip() == "help":
            print("""valid commands:
    label NAME      set label to NAME
    write FILE      write to file
    help            print this message
    exit""")
        elif cmd.strip() in ("exit", "quit", "q"):
            break
        else:
            pass

if __name__ == "__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("image", help="image to draw in the background")
    args = parser.parse_args()

    im = Image.open(args.image)
    nx, ny = im.size
    im_array = np.array(im.getdata()).reshape([ny, nx, 4])

    try:
        state = State()
        digitizer = Digitizer(state, im_array[:,:,:3]/255.0)
        plt.ion()
        plt.show()

        loop = asyncio.get_event_loop()
        loop.run_until_complete(handle_cli_cmd(state))
        loop.close()
    except KeyboardInterrupt:
        sys.exit(0)