Robotics-Mechanics/Part-8-VisualizeRobotFrames/visualize.py at main · SakshayMahna/Robotics-Mechanics · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
"""
Visualize the transformations

Matplotlib:
quiver plot
"""

from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
import numpy as np

# Function to plot a single transformation
def plot_transformation(transformation):
    """
    Plot Transformation matrix

    ...

    Parameters
    ---
    transformation: 4x4 transformation matrix

    Returns
    ---
    None

    Notes
    ---
    RGB -> XYZ
    """
    fig = plt.figure()
    ax = fig.gca(projection='3d')

    # x, y, z of 6 arrows in a quiver plot
    x = np.array([0, 0, 0, transformation[0, 3], transformation[0, 3], transformation[0, 3]])
    y = np.array([0, 0, 0, transformation[1, 3], transformation[1, 3], transformation[1, 3]])
    z = np.array([0, 0, 0, transformation[2, 3], transformation[2, 3], transformation[2, 3]])

    # u, v, w of 6 arrows in a quiver plot
    u = np.concatenate([np.array([1, 0, 0]), transformation[:3, 0]])
    v = np.concatenate([np.array([0, 1, 0]), transformation[:3, 1]])
    w = np.concatenate([np.array([0, 0, 1]), transformation[:3, 2]])

    # Color(RGB) for 6 arrows, original X, Y, Z and then transformed X, Y, Z
    red = np.array([1, 0, 0])
    green = np.array([0, 1, 0])
    blue = np.array([0, 0, 1])
    colors = np.array([red, green, blue, red, green, blue])

    q = ax.quiver(x, y, z, u, v, w, length=0.05, colors = colors, lw=1)

    plt.plot([x[0], x[3]], [y[0], y[3]], [z[0], z[3]], '--', color = 'black')

    plt.show()


# Function to plot a list of transformations
def plot_transformations(transformations):
    """
    Plot Transformation matrix

    ...

    Parameters
    ---
    transformation: list of 4x4 transformation matrix

    Returns
    ---
    None

    Notes
    ---
    RGB -> XYZ
    """
    fig = plt.figure()
    ax = fig.gca(projection='3d')

    x = np.array([])
    y = np.array([])
    z = np.array([])

    u = np.array([])
    v = np.array([])
    w = np.array([])

    red = np.array([1, 0, 0])
    green = np.array([0, 1, 0])
    blue = np.array([0, 0, 1])
    colors = []

    for transformation in transformations:
        x = np.concatenate([x, [transformation[0, 3], transformation[0, 3], transformation[0, 3]]])
        y = np.concatenate([y, [transformation[1, 3], transformation[1, 3], transformation[1, 3]]])
        z = np.concatenate([z, [transformation[2, 3], transformation[2, 3], transformation[2, 3]]])

        u = np.concatenate([u, transformation[:3, 0]])
        v = np.concatenate([v, transformation[:3, 1]])
        w = np.concatenate([w, transformation[:3, 2]])

        colors.append(red)
        colors.append(green)
        colors.append(blue)

    [0, 0, 0, 1, 1, 1]

    q = ax.quiver(x, y, z, u, v, w, length=0.05, colors = colors, lw=1)

    for i in range(x.shape[0] - 3):
        plt.plot([x[i], x[i+3]], [y[i], y[i+3]], [z[i], z[i+3]], '--', color = 'black')

    plt.show()