plot.py

   1 # -*- coding: utf-8 -*-
   2 """This scipt loads scenario and result trajectory from files and plots it."""
   3 from json import loads
   4 from math import copysign, cos, pi, sin
   5 from matplotlib import pyplot as plt
   6 from sys import argv
   7
   8 sign = lambda x: copysign(1, x)
   9
  10 HEIGHT = 1.450
  11 LENGTH = 3.760
  12 SAFETY_DIST = 0
  13 WHEEL_BASE = 2.450
  14 WIDTH = 1.625
  15
  16 SCEN_FILE = argv[1]
  17 TRAJ_FILE = argv[2]
  18 PLOT = {
  19         "enable" : True,
  20         "node" : False,
  21         "edge" : True,
  22         "sample" : False,
  23         "frame" : False,
  24         "path" : False,
  25         "traj" : True,
  26         }
  27 COLOR = {
  28         "node": "lightgrey",
  29         "edge": "lightgrey",
  30         "start": "violet",
  31         "goal": "red",
  32         "sample": "magenta",
  33         "path": "grey",
  34         "trajectory": "blue",
  35         "trajectory-frame": "lightblue",
  36         "obstacle": "black",
  37         "log": ("gold", "orange", "blueviolet", "blue", "navy", "black"),
  38         }
  39
  40 def car_frame(pose):
  41     """Return `xcoords`, `ycoords` arrays of car frame.
  42
  43     Keyword arguments:
  44     pose -- The pose of a car.
  45     """
  46     dr = (LENGTH - WHEEL_BASE) / 2
  47     df = LENGTH - dr
  48
  49     lfx = pose[0]
  50     lfx += (WIDTH / 2.0) * cos(pose[2] + pi / 2.0)
  51     lfx += df * cos(pose[2])
  52     lfx += SAFETY_DIST * cos(pose[2])
  53
  54     lrx = pose[0]
  55     lrx += (WIDTH / 2.0) * cos(pose[2] + pi / 2.0)
  56     lrx += -dr * cos(pose[2])
  57     lrx += -SAFETY_DIST * cos(pose[2])
  58
  59     rrx = pose[0]
  60     rrx += (WIDTH / 2.0) * cos(pose[2] - pi / 2.0)
  61     rrx += -dr * cos(pose[2])
  62     rrx += -SAFETY_DIST * cos(pose[2])
  63
  64     rfx = pose[0]
  65     rfx += (WIDTH / 2.0) * cos(pose[2] - pi / 2.0)
  66     rfx += df * cos(pose[2])
  67     rfx += SAFETY_DIST * cos(pose[2])
  68
  69     lfy = pose[1]
  70     lfy += (WIDTH / 2.0) * sin(pose[2] + pi / 2.0)
  71     lfy += df * sin(pose[2])
  72     lfy += SAFETY_DIST * sin(pose[2])
  73
  74     lry = pose[1]
  75     lry += (WIDTH / 2.0) * sin(pose[2] + pi / 2.0)
  76     lry += -dr * sin(pose[2])
  77     lry += -SAFETY_DIST * sin(pose[2])
  78
  79     rry = pose[1]
  80     rry += (WIDTH / 2.0) * sin(pose[2] - pi / 2.0)
  81     rry += -dr * sin(pose[2])
  82     rry += -SAFETY_DIST * sin(pose[2])
  83
  84     rfy = pose[1]
  85     rfy += (WIDTH / 2.0) * sin(pose[2] - pi / 2.0)
  86     rfy += df * sin(pose[2])
  87     rfy += SAFETY_DIST * sin(pose[2])
  88
  89     xcoords = (lfx, lrx, rrx, rfx)
  90     ycoords = (lfy, lry, rry, rfy)
  91     return (xcoords, ycoords)
  92
  93 def load_scenario(fname):
  94     """Load scenario from file."""
  95     if fname is None:
  96         raise ValueError("File name as argument needed")
  97     with open(fname, "r") as f:
  98         scenario = loads(f.read())
  99     return scenario
 100
 101 def load_trajectory(fname):
 102     """Load trajectory from file."""
 103     if fname is None:
 104         raise ValueError("File name as argument needed")
 105     with open(fname, "r") as f:
 106         trajectory = loads(f.read())
 107         return trajectory
 108
 109 def plot_nodes(nodes=[]):
 110     """Return xcoords, ycoords of nodes to plot.
 111
 112     Keyword arguments:
 113     nodes -- The list of nodes to plot.
 114     """
 115     xcoords = []
 116     ycoords = []
 117     for n in nodes:
 118         xcoords.append(n[0])
 119         ycoords.append(n[1])
 120     return (xcoords, ycoords)
 121
 122 def plot_segments(segments=[]):
 123     """Return xcoords, ycoords of segments to plot.
 124
 125     Keyword arguments:
 126     segments -- The list of segments to plot.
 127     """
 128     pass
 129
 130 if __name__ == "__main__":
 131     s = load_scenario(SCEN_FILE)
 132     try:
 133         t = load_trajectory(TRAJ_FILE) # fixed to trajectories
 134     except:
 135         pass
 136
 137     plt.rcParams["font.size"] = 24
 138     fig = plt.figure()
 139     ## 1st subplot
 140     ax = fig.add_subplot(121)
 141     ax.set_aspect("equal")
 142     ax.set_title("Final path")
 143     ax.set_xlabel("x [m]")
 144     ax.set_ylabel("y [m]")
 145
 146     # plot here
 147     for o in s["obst"]:
 148         try:
 149             plt.plot(*plot_nodes(o["segment"]), color="black")
 150         except:
 151             pass
 152         try:
 153             ax.add_artist(plt.Circle((o["circle"][0], o["circle"][1]),
 154                     o["circle"][2],
 155                     color="black", fill=False, hatch="/"))
 156         except:
 157             pass
 158     if PLOT["node"]:
 159         try:
 160             plt.plot(*plot_nodes(t["node"]), color=COLOR["node"],
 161                     marker=".", linestyle = "None")
 162         except:
 163             print("No RRTNode")
 164     if False:
 165         try:
 166             for edges in t["edge"]:
 167                 for e in edges:
 168                     plt.plot([e[0][0], e[1][0]], [e[0][1], e[1][1]],
 169                             color=COLOR["edge"])
 170         except:
 171             print("No edges")
 172     if PLOT["sample"]:
 173         try:
 174             if PLOT["frame"]:
 175                 for i in t["samp"]:
 176                     plt.plot(*car_frame(i), color=COLOR["sample"])
 177             else:
 178                 plt.plot(*plot_nodes(t["samp"]), color=COLOR["sample"],
 179                         marker=".", linestyle = "None")
 180         except:
 181             print("No RRTSample")
 182     if PLOT["path"]:
 183         try:
 184             for path in range(len(t["path"])):
 185                 plt.plot(*plot_nodes(t["path"][path]), color=COLOR["path"])
 186         except:
 187             print("No path")
 188     if PLOT["traj"]:
 189         try:
 190             for traj in range(len(t["traj"])):
 191                 if PLOT["frame"]:
 192                     for i in t["traj"][traj]:
 193                         plt.plot(*car_frame(i), color=COLOR["log"][traj],
 194                                 label=t["cost"][traj])
 195                 else:
 196                     try:
 197                         plt.plot(
 198                                 *plot_nodes(t["traj"][traj]),
 199                                 color=COLOR["log"][traj],
 200                                 label=t["cost"][traj])
 201                     except:
 202                         plt.plot(
 203                                 *plot_nodes(t["traj"][traj]),
 204                                 color="black",
 205                                 label=t["cost"][traj])
 206         except:
 207             print("No trajectory")
 208     plt.plot(*car_frame(s["init"]), color="red", lw=2)
 209     plt.plot(*car_frame(s["goal"]), color="red", lw=2)
 210     plt.plot(*plot_nodes([s["init"]]), color="red", marker="+", ms=12)
 211     plt.plot(*plot_nodes([s["goal"]]), color="red", marker="+", ms=12)
 212     # end plot here
 213
 214     handles, labels = ax.get_legend_handles_labels()
 215     #lgd = ax.legend(handles, labels, loc="upper center",
 216     #        bbox_to_anchor=(0.5, -0.11), title="Cost")
 217     ## 2nd subplot
 218     ax = fig.add_subplot(122)
 219     ax.set_aspect("equal")
 220     ax.set_title("All edges")
 221     ax.set_xlabel("x [m]")
 222     ax.set_ylabel("y [m]")
 223
 224     # plot here
 225     for o in s["obst"]:
 226         try:
 227             plt.plot(*plot_nodes(o["segment"]), color="black")
 228         except:
 229             pass
 230         try:
 231             ax.add_artist(plt.Circle((o["circle"][0], o["circle"][1]),
 232                 o["circle"][2],
 233                 color="black", fill=False))
 234         except:
 235             pass
 236     if PLOT["node"]:
 237         try:
 238             plt.plot(*plot_nodes(t["node"]), color=COLOR["node"],
 239                     marker=".", linestyle = "None")
 240         except:
 241             print("No RRTNode")
 242     if True:
 243         try:
 244             for edges in t["edge"]:
 245                 for e in edges:
 246                     plt.plot([e[0][0], e[1][0]], [e[0][1], e[1][1]],
 247                             color=COLOR["edge"])
 248         except:
 249             print("No edges")
 250     if PLOT["sample"]:
 251         try:
 252             if PLOT["frame"]:
 253                 for i in t["samp"]:
 254                     plt.plot(*car_frame(i), color=COLOR["sample"])
 255             else:
 256                 plt.plot(*plot_nodes(t["samp"]), color=COLOR["sample"],
 257                         marker=".", linestyle = "None")
 258         except:
 259             print("No RRTSample")
 260     if PLOT["path"]:
 261         try:
 262             for path in range(len(t["path"])):
 263                 plt.plot(*plot_nodes(t["path"][path]), color=COLOR["path"])
 264         except:
 265             print("No path")
 266     if False:
 267         try:
 268             for traj in range(len(t["traj"])):
 269                 if PLOT["frame"]:
 270                     for i in t["traj"][traj]:
 271                         plt.plot(*car_frame(i), color=COLOR["log"][traj],
 272                                 label=t["cost"][traj])
 273                 else:
 274                     try:
 275                         plt.plot(
 276                                 *plot_nodes(t["traj"][traj]),
 277                                 color=COLOR["log"][traj],
 278                                 label=t["cost"][traj])
 279                     except:
 280                         plt.plot(
 281                                 *plot_nodes(t["traj"][traj]),
 282                                 color="black",
 283                                 label=t["cost"][traj])
 284         except:
 285             print("No trajectory")
 286     plt.plot(*car_frame(s["init"]), color="red", lw=2)
 287     plt.plot(*car_frame(s["goal"]), color="red", lw=2)
 288     plt.plot(*plot_nodes([s["init"]]), color="red", marker="+", ms=12)
 289     plt.plot(*plot_nodes([s["goal"]]), color="red", marker="+", ms=12)
 290     # end plot here
 291
 292     handles, labels = ax.get_legend_handles_labels()
 293     #lgd = ax.legend(handles, labels, loc="upper center",
 294     #        bbox_to_anchor=(0.5, -0.11), title="Cost")
 295
 296     # END OF SUBPLOTS
 297     plt.show()
 298     #plt.savefig("{}.png".format(argv[2]), bbox_extra_artists=(lgd,),
 299     #        bbox_inches='tight')
 300     plt.close(fig)