AE4441P-OperationsOptimisations/plotting.py at main · JoshuaBog09/AE4441P-OperationsOptimisations · 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
113
114
115
116
117
118
119
120
import matplotlib.pyplot as plt
from matplotlib.patches import Rectangle
from matplotlib.animation import FuncAnimation, PillowWriter
from receding_horizon import Scene, Vehicle, Config, Results
from receding_horizon import read_setup
import numpy as np
import pickle
import sys
import os


def static_plot(scene: Scene, results: Results):
    # Plot the scene
    fig, ax = plt.subplots()

    # Set limits
    ax.set_xlim(scene.map_bounds[0, 0], scene.map_bounds[1, 0])
    ax.set_ylim(scene.map_bounds[0, 1], scene.map_bounds[1, 1])

    # Plot vehicle location
    # ax.plot(x_path, y_path, marker=".", color='red', label="Path")
    ax.plot(scene.goal[0], scene.goal[1], marker="X", color='green', label="Goal")

    # Plot obstacles
    for obstacle in scene.obstacles:
        origin = obstacle[0]
        delta = obstacle[1] - obstacle[0]
        width = delta[0]
        height = delta[1]

        ax.add_patch(Rectangle(origin, width, height, color='dimgrey'))

    for path in results.plan_path:
        ax.plot(path[:, 0], path[:, 1], marker="", color='blue', linewidth=0.33, zorder=0)

    results.exec_path = np.concatenate(results.exec_path)

    ax.plot(results.exec_path[:, 0], results.exec_path[:, 1], marker=".", color='red', linewidth=1, zorder=10,
            label="Executed Path", markersize=5.5, markevery=3)

    # display plot
    plt.legend()
    plt.show()


def animate(i, ax, results: Results, config: Config):
    ax.plot(results.plan_path[i][:, 0], results.plan_path[i][:, 1], marker="", color='blue', linewidth=0.33, zorder=0)

    if i == 0:
        ax.plot(results.exec_path[i][:, 0], results.exec_path[i][:, 1], marker=".", color='red', linewidth=1, zorder=10)
    else:
        path = np.vstack((results.exec_path[i-1][-1, :], results.exec_path[i]))
        ax.plot(path[:, 0], path[:, 1], marker=".", color='red', linewidth=1, zorder=10)


def animated_plot(scene: Scene, results: Results, config: Config):
    fig, ax = plt.subplots()
    # Set limits
    ax.set_xlim(scene.map_bounds[0, 0], scene.map_bounds[1, 0])
    ax.set_ylim(scene.map_bounds[0, 1], scene.map_bounds[1, 1])

    # Plot vehicle location
    # ax.plot(x_path, y_path, marker=".", color='red', label="Path")
    ax.plot(scene.goal[0], scene.goal[1], marker="X", color='green', label="Goal")

    # Plot obstacles
    for obstacle in scene.obstacles:
        origin = obstacle[0]
        delta = obstacle[1] - obstacle[0]
        width = delta[0]
        height = delta[1]

        ax.add_patch(Rectangle(origin, width, height, color='dimgrey'))
    ani = FuncAnimation(fig, animate, fargs=(ax, results, config), frames=len(results.exec_path))

    major_ticks_x = np.arange(scene.map_bounds[0, 0], scene.map_bounds[1, 0], 10)
    minor_ticks_x = np.arange(scene.map_bounds[0, 0], scene.map_bounds[1, 0], 1)

    major_ticks_y = np.arange(scene.map_bounds[0, 1], scene.map_bounds[1, 1], 10)
    minor_ticks_y = np.arange(scene.map_bounds[0, 1], scene.map_bounds[1, 1], 1)

    ax.set_xticks(major_ticks_x)
    ax.set_yticks(major_ticks_y)
    ax.set_xticks(minor_ticks_x, minor=True)
    ax.set_yticks(minor_ticks_y, minor=True)

    ax.grid(which="major", alpha=0.6)
    ax.grid(which="minor", alpha=0.3)

    ax.set_aspect("equal")

    plt.legend()
    plt.show()

    return ani


if __name__ == "__main__":

    if len(sys.argv) < 2:
        print("Directory not provided as input argument")
        raise FileNotFoundError

    location = sys.argv[1]

    setup_file = os.path.join(location, "setup.yaml")
    results_file = os.path.join(location, "result.pickle")
    plot_file = os.path.join(location, "result.gif")

    map_data, config_data, vehicle_data = read_setup(setup_file)

    map     = Scene.from_dict(map_data)
    config  = Config.from_dict(config_data)
    vehicle = Vehicle.from_dict(vehicle_data)

    with open(results_file, 'rb') as f:
        results = pickle.load(f)

    ani = animated_plot(map, results, config)
    ani.save(plot_file, writer=PillowWriter(fps=30))