B Spline Pure Numpy

bspline.py

import numpy as np
import matplotlib.patches as patches


class BSpline():
    def __init__(self):
        pass
        self.warp = 4.
        self.count = 26*2


    def process(self, Px, Py, Pw, ax, MODE=1):
        self.mode = MODE # Can be 0 or 1 for terminal tangents
        self.ax = ax
        self.Px = Px
        self.Py = Py
        self.Pw = Pw

        count = self.count
        if self.mode == 0:
            self.n = len(Px) - 2
        else:
            self.n = len(Px) - 2
        self.n1 = self.n+1;
        self.B0 = [None]*count
        self.B1 = [None]*count
        self.B2 = [None]*count
        self.B3 = [None]*count
        self.dx = [None]*self.n
        self.dy = [None]*self.n

        t = 0.;
        for i in range(0, count):
            t1 = 1-t
            t12 = t1*t1
            t2 = t*t
            self.B0[i] = t1*t12
            self.B1[i] = 3*t*t12
            self.B2[i] = 3*t2*t1
            self.B3[i] = t*t2
            t += 1.00/count

        return self.drawSpline()

    def findCPoints(self):
        if self.mode == 0:
            self.dx[0] = self.Px[self.n] - self.Px[0];  
            self.dy[0] = self.Py[self.n] - self.Py[0];
            self.dx[self.n-1] = -(self.Px[self.n1] - self.Px[self.n-1])
            self.dy[self.n-1] = -(self.Py[self.n1] - self.Py[self.n - 1])
        else:
            DIV = 3.
            self.dx[0] = (self.Px[1] - self.Px[0])/DIV;  
            self.dy[0] = (self.Py[1] - self.Py[0])/DIV;
            self.dx[self.n-1] = (self.Px[self.n-1] - self.Px[self.n-2])/DIV;  
            self.dy[self.n-1] = (self.Py[self.n-1] - self.Py[self.n-2])/DIV;

        #self.warp += 0.05

        warps = self.Pw

        Ax = [None]*self.n
        Ay = [None]*self.n
        Bi = [None]*self.n
        Bi[1] = -1/warps[1]
        Ax[1] = -(self.Px[2] - self.Px[0] - self.dx[0])*Bi[1]
        Ay[1] = -(self.Py[2] - self.Py[0] - self.dy[0])*Bi[1];
        for i in range(2, self.n-1):
            warpval = warps[i]
            Bi[i] = -1/(warpval+ Bi[i-1]);
            Ax[i] = -(self.Px[i+1] - self.Px[i-1] - Ax[i-1])*Bi[i];
            Ay[i] = -(self.Py[i+1] - self.Py[i-1] - Ay[i-1])*Bi[i];
        for i in range(self.n-2, 0, -1):
            self.dx[i] = Ax[i] + self.dx[i+1]*Bi[i];  
            self.dy[i] = Ay[i] + self.dy[i+1]*Bi[i];

        #self.dx[1] = 0
        #self.dy[0] = 0
        #self.dy[1] = 0
        # for i in range(0, len(self.dx)):
        #     self.dx[i] = 0
        #     self.dy[i] = 0
        # print self.dx
        # print self.dy


    def drawSpline(self):
        self.findCPoints()

        w = 1.
        h = 1.
        d = 0.
        h1 = h
        d2 = d
        step = 1./w
        t = step;
        scPx = [None]*self.n
        scPy = [None]*self.n
        scDx = [None]*self.n
        scDy = [None]*self.n
        X = None
        Y = None

        squaresize = 0.3
        for i in range(0, self.n):
           X = scPx[i] = self.Px[i]*w
           Y = scPy[i] = self.Py[i]*h; 
           scDx[i] = self.dx[i]*w;   
           scDy[i] = self.dy[i]*h;
           rect = patches.Rectangle((X - squaresize/2.,Y - squaresize/2.),squaresize,squaresize,linewidth=1,edgecolor='r',facecolor='none')
           self.ax.add_patch(rect)

        #if self.mode == 0:
        self.ax.add_patch(patches.Rectangle((scPx[0]+scDx[0] - squaresize/2., (scPy[0]+scDy[0])- squaresize/2.),squaresize,squaresize,
            linewidth=1,edgecolor='b',facecolor='none'))
        self.ax.add_patch(patches.Rectangle((scPx[self.n-1]-scDx[self.n-1] - squaresize/2., (scPy[self.n-1]-scDy[self.n-1])- squaresize/2.),squaresize,squaresize,
            linewidth=1,edgecolor='b',facecolor='none'))

        # if self.n > 1:
        #     paths = [[scPx[0], scPy[0]]]
        #     for i in range(1, self.n):
        #         paths.append([scPx[i-1]+scDx[i-1], scPy[i-1]-scDy[i-1]])
        #         paths.append([scPx[i]-scDx[i], scPy[i]+scDy[i]])
        #         paths.append([scPx[i], scPy[i]])
        #     paths_arr = np.array(paths)
        #     plt.plot(paths_arr[:, 0], paths_arr[:, 1])
        #     #for path in paths:

        paths = [[scPx[0], scPy[0]]]
        for i in range(0, self.n-1):
            for k in range(0, self.count):
                X = (scPx[i]*self.B0[k] + (scPx[i] + scDx[i])*self.B1[k] +
                    (scPx[i+1] - scDx[i+1])*self.B2[k] + scPx[i+1]*self.B3[k])
                Y = (scPy[i]*self.B0[k] + (scPy[i] + scDy[i])*self.B1[k] +
                    (scPy[i+1] - scDy[i+1])*self.B2[k] + scPy[i+1]*self.B3[k]);
                paths.append([X,Y])
        paths_arr = np.array(paths)
        plt.scatter(paths_arr[:, 0], paths_arr[:, 1], s=0.2)

        return paths_arr

# cv = np.array(
# [[ 2.70967742,  1.30411255,  4.        ], # D
#  [ 2.45564516,  6.28246753,  4.        ], # C
#  [-0.90322581,  6.39069264,  4.        ],
#  [-1.04435484,  1.27705628,  4.        ],
#  [-5.75806452,  1.41233766,  4.        ], # B
#  [-5.95564516,  6.14718615,  4.        ], # A
#  [ 2.51209677,  3.79329004,  4.        ], # T: Should be center of CD
#  [-6.06854839,  4.11796537,  4.        ]] # T: Should be center of AB
#     )

cv = np.array(
[[ 6.35080645,  5.95779221,  4.        ],
 [ 6.32258065,  1.46645022,  4.        ],
 [ 4.20564516,  1.46645022,  4.        ],
 [ 4.20564516,  6.01190476,  4.        ],
 [ 1.80645161,  5.95779221,  4.        ],
 [ 1.66532258,  1.46645022,  4.        ],
 [-0.56451613,  1.38528139,  4.        ],
 [-0.64919355,  6.06601732,  4.        ],
 [-2.73790323,  6.03896104,  4.        ],
 [-2.96370968,  1.33116883,  4.        ],
 [-5.75806452,  1.41233766,  4.        ],
 [-5.95564516,  6.14718615,  4.        ],
 [ 6.29435484,  3.76623377,  4.        ],
 [-6.06854839,  4.11796537,  4.        ]]
    )

class Click():
    def __init__(self, ax, button=1):
        self.ax=ax
        self.button=button
        self.press=False
        self.move = False
        self.c1=self.ax.figure.canvas.mpl_connect('button_press_event', self.onpress)
        self.c2=self.ax.figure.canvas.mpl_connect('button_release_event', self.onrelease)
        self.c3=self.ax.figure.canvas.mpl_connect('motion_notify_event', self.onmove)

        self.c4=self.ax.figure.canvas.mpl_connect('scroll_event',self.onzoom)

    def onzoom(self, event):
        print("zoom")
        global cv
        lowest_dist = 1000
        lowest_index = -1
        for i in range(0, cv.shape[0]):
            dist = np.sqrt((cv[i,0]-event.xdata)**2 + (cv[i,1]-event.ydata)**2)
            if dist < lowest_dist:
                lowest_dist = dist
                lowest_index = i
        if lowest_dist < 2:
            if event.button == "up":
                cv[lowest_index, 2] += 0.5
            elif event.button == "down":
                if cv[lowest_index, 2] > 1.5:
                    cv[lowest_index, 2] -= 0.5

    def onclick(self,event):
        if event.inaxes == self.ax:
            global cv
            lowest_dist = 1000
            lowest_index = -1
            for i in range(0, cv.shape[0]):
                dist = np.sqrt((cv[i,0]-event.xdata)**2 + (cv[i,1]-event.ydata)**2)
                if dist < lowest_dist:
                    lowest_dist = dist
                    lowest_index = i
            if event.button == self.button:
                if lowest_dist > 1:
                    cv = np.vstack((np.array([event.xdata, event.ydata, 4.]),cv))

                    # Get Centre
                    print(cv[0,:])
                    print(cv[1,:])
                    center_x = (cv[0,0] + cv[1,0])/2
                    center_y = (cv[0,1] + cv[1,1])/2
                    print(center_x, center_y)
                    cv[-2,0] = center_x
                    cv[-2,1] = center_y
            if event.button == 3:

                if lowest_dist < 1:
                    cv = np.delete(cv, (lowest_index), axis=0)

    def onpress(self,event):
        self.press=True
    def onmove(self,event):
        if self.press:
            self.move=True

            # Find the closest point to the click
            global cv
            lowest_dist = 1000
            lowest_index = -1
            for i in range(0, cv.shape[0]):
                dist = np.sqrt((cv[i,0]-event.xdata)**2 + (cv[i,1]-event.ydata)**2)
                if dist < lowest_dist:
                    lowest_dist = dist
                    lowest_index = i
            if lowest_dist < 2:
                cv[lowest_index, 0] = event.xdata
                cv[lowest_index, 1] = event.ydata
                print(lowest_index)

            print(cv)

    def onrelease(self,event):
        if self.press and not self.move:
            self.onclick(event)
        self.press=False; self.move=False
        #print("Release")

if __name__ == "__main__":
    import matplotlib.pyplot as plt
    first_run = True
    click  = None

    axes = plt.gca()
    
    bspline = BSpline()

    def test(closed):
        global click
        global first_run
        global cv

        cvT = cv.T
        ax2 = plt.gca()
        spline = bspline.process(cvT[0],cvT[1],cvT[2],ax2)
        # plt.plot(x,y,'k-',label='Curve')
        plt.minorticks_on()
        plt.legend()
        plt.xlabel('x')
        plt.ylabel('y')
        plt.xlim(-7, 7)
        plt.ylim(0, 10)
        #plt.gca().set_aspect('equal', adjustable='box-forced')
        plt.pause(0.005)
        plt.cla()

        if(first_run):
            first_run = False
            click = Click(ax2, button=1)

    while 1:
        test(False)