대화식으로 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)



기능


  • 왼쪽 클릭을 누른 상태에서 궤적을 그립니다
  • 오른쪽 클릭으로 지우기
  • 키보드의 $ r $를 누르면 pickle로 궤적을 저장합니다.
  • 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()
    
    

    좋은 웹페이지 즐겨찾기