インタラクティブにmatplotlibでお絵かき

はじめに

matplotlibでインタラクティブにお絵かきして，その軌跡を保存したかったので作ってみました

どんなものか

FuncAnimationを使って連続して描画しています．
そして，描画中にcallback関数を呼ぶことでマウス，キーボードのイベントを取得します．
matplotlibで準備されているcanvas.mpl_connectを使用します．

fig.canvas.mpl_connect('motion_notify_event', self.onclick)
fig.canvas.mpl_connect('key_press_event', self.onkey)

機能

1. 左クリックを押しながら軌跡を描きます
2. 右クリックで消去
3. キーボードの$r$を押すとpickleで軌跡を保存します．
4. LineDrawer.check_traj()で保存した軌跡を確認することができます．

コード

import math
import pickle
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import animation


class LineDrawer():
    def __init__(self):
        self.trajectory = None
        self.xx = []
        self.yy = []


    def run(self):
        #--------- prepare a figure
        fig = plt.figure()
        ax = fig.add_subplot(1,1,1, xlim=(0,300), ylim=(-50,50))
        ax.xaxis.grid(); ax.yaxis.grid()  # Grid
        
        #--------- hook up function to mouse event
        self.traj_line, = ax.plot([], [], 'o', markersize=4, mew=4)
        self.record_line, = ax.plot([], [], 'o', markersize=4, mew=4, color='m')
        self.empty, = ax.plot([],[])

        fig.canvas.mpl_connect('motion_notify_event', self.onclick)
        fig.canvas.mpl_connect('key_press_event', self.onkey)

        #--------- run animation 
        anim = animation.FuncAnimation(fig, self.anim_animate,
                init_func=self.anim_init, frames=30, interval=100, blit=True)
        self.anim = anim

    def onclick(self, event):
        # get mouse position and scale appropriately to convert to (x,y)
        if event.xdata is not None:
            self.trajectory = np.array([event.xdata, event.ydata])
            print('event.x=%d, event.y=%d, event.xdata=%f, event.ydata=%f'\
                 %(event.x, event.y, event.xdata, event.ydata))
        if event.button==1:
            self.xx.append(event.xdata)
            self.yy.append(event.ydata)
            print('event.button=%d,  event.x=%d, event.y=%d, event.xdata=%f, event.ydata=%f'\
                 %(event.button, event.x, event.y, event.xdata, event.ydata))
        if event.button==3:
            self.xx = []
            self.yy = []
                    

    def onkey(self, event):
        # Record
        if event.key == 'r':
            traj = np.array([self.xx, self.yy])
            with open('traj.pickle', 'wb') as f:
                pickle.dump(traj, f)


    def anim_init(self):
        self.traj_line.set_data([], [])
        self.record_line.set_data([], [])
        self.empty.set_data([], [])

        return self.traj_line, self.record_line, self.empty


    def anim_animate(self, i):
        if self.trajectory is not None:
            self.traj_line.set_data(self.trajectory)

        if self.xx is not None:
            self.record_line.set_data(self.xx, self.yy)
        
        self.empty.set_data([],[])
        
        return self.traj_line, self.record_line, self.empty
        
    
    def show(self):
        try:
            plt.show()
        except AttributeError:
            pass

    def check_traj(self):
        with open('traj.pickle', 'rb') as f:
            traj = pickle.load(f)

        fig = plt.figure()
        ax = fig.add_subplot(111, xlim=(0,300), ylim=(-50,50))
        ax.grid(True)
        ax.plot(traj[0], traj[1], 'o', markersize=4, mew=4, color='k')

if __name__=='__main__':
    ld = LineDrawer()
    ld.run()
    ld.show()
    ld.check_traj()
    ld.show()