Interactive Bezier curve builder for jupyter notebook

interactive_bezier_curves.py

%matplotlib widget # need to install ipympl

import numpy as np
from matplotlib.lines import Line2D
from matplotlib.artist import Artist
import matplotlib
import matplotlib.pyplot as plt
from matplotlib.patches import Polygon

import numpy as np
from scipy.special import binom

import matplotlib.pyplot as plt

from matplotlib.lines import Line2D
from ipywidgets import widgets
from IPython.display import display


# Code modified from: https://gist.github.com/astrojuanlu/7284462 
# and https://matplotlib.org/stable/gallery/event_handling/poly_editor.html
class BezierInteractor:
    """
    An interactive bezier curve builder editor.

    Key-bindings

      't' toggle vertex markers on and off.  When vertex markers are on,
          you can move them, delete them

      'd' delete the vertex under point

      'i' insert a vertex at point.  You must be within epsilon of the
          line connecting two existing vertices
          
    You can also draging vertexes with the mouse to modify them.
    """

    showverts = True
    epsilon = 0.1  # max distance to count as a vertex hit

    def __init__(self, poly):            
        self.poly = poly
        self.xp = list(poly.get_xdata())
        self.yp = list(poly.get_ydata())
        self.ax = poly.axes
        canvas = poly.figure.canvas
        
        # Event handler for mouse clicking
        canvas.mpl_connect('draw_event', self.on_draw)
        canvas.mpl_connect('button_press_event', self.on_button_press)
        canvas.mpl_connect('key_press_event', self.on_key_press)
        canvas.mpl_connect('button_release_event', self.on_button_release)
        canvas.mpl_connect('motion_notify_event', self.on_mouse_move)
        self.canvas = canvas

        # Create Bézier curve
        line_bezier = Line2D([], [],
                             c=poly.get_markeredgecolor())
        self.bezier_curve = self.ax.add_line(line_bezier)
        

        self.cid = self.poly.add_callback(self.poly_changed)
        self._ind = None  # the active vert

    def on_draw(self, event):
        self.background = self.canvas.copy_from_bbox(self.ax.bbox)
        self.ax.draw_artist(self.poly)

    def poly_changed(self, poly):
        """This method is called whenever the pathpatch object is called."""
        # only copy the artist props to the line (except visibility)
        vis = self.bezier_curve.get_visible()
        Artist.update_from(self.bezier_curve, poly)
        self.bezier_curve.set_visible(vis)  # don't use the poly visibility state

    def get_ind_under_point(self, event):
        """
        Return the index of the point closest to the event position or *None*
        if no point is within ``self.epsilon`` to the event position.
        """
        d = np.hypot(np.array(self.xp) - event.xdata, np.array(self.yp) - event.ydata)
        indseq, = np.nonzero(d == d.min())
        ind = indseq[0]

        if d[ind] >= self.epsilon:
            ind = None

        return ind

    def on_button_press(self, event):
        """Callback for mouse button presses."""
        if not self.showverts:
            return
        if event.inaxes is None:
            return
        if event.button != 1:
            return
        self._ind = self.get_ind_under_point(event)

    def on_button_release(self, event):
        """Callback for mouse button releases."""
        if not self.showverts:
            return
        if event.button != 1:
            return
        self._ind = None

    def on_key_press(self, event):
        """Callback for key presses."""
        if not event.inaxes:
            return
        if event.key == 't':
            self.showverts = not self.showverts
            self.poly.set_visible(self.showverts)
            if not self.showverts:
                self._ind = None
        elif event.key == 'd':
            ind = self.get_ind_under_point(event)
            if ind is not None:
                self.xp.pop(ind)
                self.yp.pop(ind)
                self._update_bezier()
        elif event.key == 'i':
            self.xp.append(event.xdata)
            self.yp.append(event.ydata)
            self._update_bezier()
            
        if self.poly.stale:
            self.canvas.draw_idle()
            
    def on_mouse_move(self, event):
        """Callback for mouse movements."""
        if not self.showverts:
            return
        if self._ind is None:
            return
        if event.inaxes is None:
            return
        if event.button != 1:
            return

        self.xp[self._ind] = event.xdata
        self.yp[self._ind] = event.ydata
        self._update_bezier()
            
    def _build_bezier(self):
        x, y = Bezier(list(zip(self.xp, self.yp))).T
        return x, y

    def _update_bezier(self):
        self.poly.set_data(self.xp, self.yp)
        self.bezier_curve.set_data(*self._build_bezier())
        self.canvas.draw()


def Bernstein(n, k):
    """Bernstein polynomial.
    """
    coeff = binom(n, k)

    def _bpoly(x):
        return coeff * x ** k * (1 - x) ** (n - k)

    return _bpoly


def Bezier(points, num=200):
    """Build Bézier curve from points.
    """
    N = len(points)
    t = np.linspace(0, 1, num=num)
    curve = np.zeros((num, 2))
    for ii in range(N):
        curve += np.outer(Bernstein(N - 1, ii)(t), points[ii])
    return curve


# Initial setup
    
dpi=200
fig, ax = plt.subplots(1, 1, figsize=(1920/dpi, 1080/dpi))

# Empty line
line = Line2D([], [], ls='--', c='#666666',
              marker='x', mew=2, mec='#204a87')
ax.add_line(line)

# Canvas limits
ax.set_xlim(0, 1)
ax.set_ylim(0, 1)
ax.set_title("Bézier curve")

# Create BezierBuilder
bezier_builder = BezierInteractor(line)
plt.show()


def reset_plot(*args, **kwargs):
    """Reset the interactive plots to inital values."""
    bezier_builder.xp=[]
    bezier_builder.yp=[]
    bezier_builder._update_bezier()
    

reset_plot()
reset_button = widgets.Button(description = "Reset")
reset_button.on_click(reset_plot)
display(reset_button)