plot.py 2.46 KB
Newer Older
Guilhem Saurel's avatar
Guilhem Saurel committed
1
import eigenpy
stevet's avatar
stevet committed
2
import matplotlib.pyplot as plt
stevet's avatar
stevet committed
3
from mpl_toolkits.mplot3d import Axes3D
stevet's avatar
stevet committed
4
import numpy as np
Guilhem Saurel's avatar
Guilhem Saurel committed
5
from numpy import array
stevet's avatar
stevet committed
6

Guilhem Saurel's avatar
Guilhem Saurel committed
7
from curves import bezier
stevet's avatar
stevet committed
8

Guilhem Saurel's avatar
Guilhem Saurel committed
9
10
11
12
eigenpy.switchToNumpyArray()


def plotControlPoints2D(bez, axes=[0, 1], color="r", ax=None):
stevet's avatar
stevet committed
13
14
15
16
    wps = [bez.waypointAtIndex(i) for i in range(bez.nbWaypoints)]
    x = np.array([wp[axes[0]] for wp in wps])
    y = np.array([wp[axes[1]] for wp in wps])
    if ax is not None:
Guilhem Saurel's avatar
Guilhem Saurel committed
17
18
19
20
21
22
23
24
        ax.scatter(x, y, color=color)
    else:
        plt.scatter(x, y, color=color)


def plotBezier2D(bez, axes=[0, 1], step=100., color="b", showControlPoints=False, ax=None):
    points1 = np.array([(bez(i / step * (bez.max() - bez.min()) + bez.min())[axes[0]],
                         bez(i / step * (bez.max() - bez.min()) + bez.min())[axes[1]]) for i in range(int(step) + 1)])
stevet's avatar
stevet committed
25
26
    x = points1[:, 0]
    y = points1[:, 1]
Guilhem Saurel's avatar
Guilhem Saurel committed
27
    if ax is None:
28
29
        fig = plt.figure()
        ax = fig.add_subplot(111)
Guilhem Saurel's avatar
Guilhem Saurel committed
30
    ax.plot(x, y, color, linewidth=2.0)
31
    if showControlPoints:
Guilhem Saurel's avatar
Guilhem Saurel committed
32
        plotControlPoints2D(bez, color=color, ax=ax)
stevet's avatar
stevet committed
33

34

Guilhem Saurel's avatar
Guilhem Saurel committed
35
def plotControlPoints(bez, color="r", ax=None):
stevet's avatar
stevet committed
36
37
38
39
    wps = [bez.waypointAtIndex(i) for i in range(bez.nbWaypoints)]
    x = np.array([wp[0] for wp in wps])
    y = np.array([wp[1] for wp in wps])
    z = np.array([wp[2] for wp in wps])
Guilhem Saurel's avatar
Guilhem Saurel committed
40
    if ax is None:
41
42
        fig = plt.figure()
        ax = fig.add_subplot(111, projection="3d")
Guilhem Saurel's avatar
Guilhem Saurel committed
43
44
45
46
47
48
49
    ax.scatter(x, y, z, color=color)


def plotBezier(bez, step=100., color="b", linewidth=2., showControlPoints=False, ax=None):
    points1 = np.array([(bez(i / step * (bez.max() - bez.min()) + bez.min())[0],
                         bez(i / step * (bez.max() - bez.min()) + bez.min())[1],
                         bez(i / step * (bez.max() - bez.min()) + bez.min())[2]) for i in range(int(step) + 1)])
50
51
52
    x = points1[:, 0]
    y = points1[:, 1]
    z = points1[:, 2]
Guilhem Saurel's avatar
Guilhem Saurel committed
53
54

    if ax is None:
55
56
        fig = plt.figure()
        ax = fig.add_subplot(111, projection="3d")
Guilhem Saurel's avatar
Guilhem Saurel committed
57
    ax.plot(x, y, z, color, linewidth=linewidth)
58
    if showControlPoints:
Guilhem Saurel's avatar
Guilhem Saurel committed
59
        plotControlPoints(bez, color=color, ax=ax)
stevet's avatar
stevet committed
60
61
62
63
64


if __name__ == '__main__':
    waypoints = array([[1., 2., 3.], [-4., -5., -6.], [4., 5., 6.], [7., 8., 9.]]).transpose()
    a = bezier(waypoints)
Guilhem Saurel's avatar
Guilhem Saurel committed
65

stevet's avatar
stevet committed
66
    fig = plt.figure()
Guilhem Saurel's avatar
Guilhem Saurel committed
67
68
69
70
71
72
    ax = fig.add_subplot(121)
    plotBezier2D(a, axes=[1, 0], ax=ax)
    plotControlPoints2D(a, axes=[1, 0], ax=ax)
    ax = fig.add_subplot(122, projection="3d")
    plotBezier(a, ax=ax)
    plotControlPoints(a, ax=ax)
stevet's avatar
stevet committed
73
    plt.show(block=False)