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, exit
   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 PLOT = {
  17         "enable" : True,
  18         "node" : False,
  19         "edge" : True,
  20         "sample" : False,
  21         "frame" : False,
  22         "path" : False,
  23         "traj" : True,
  24         "slot" : True,
  25         }
  26 COLOR = {
  27         "node": "lightgrey",
  28         "edge": "lightgrey",
  29         "start": "violet",
  30         "goal": "red",
  31         "sample": "magenta",
  32         "path": "grey",
  33         "trajectory": "blue",
  34         "trajectory-frame": "lightblue",
  35         "obstacle": "black",
  36         "log": ("gold", "orange", "blueviolet", "blue", "navy", "black"),
  37         "slot": "red"
  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     SCEN_FILE = argv[1]
 132
 133     s = load_scenario(SCEN_FILE)
 134     try:
 135         TRAJ_FILE = argv[2]
 136         t = load_trajectory(TRAJ_FILE) # fixed to trajectories
 137     except:
 138         t = {"elap": 0}
 139
 140     plt.rcParams["font.size"] = 24
 141     plt.rcParams['hatch.linewidth'] = 1.0
 142     fig = plt.figure()
 143
 144     ################
 145     ## 1st subplot
 146
 147     if "edge" in t.keys():
 148         ax = fig.add_subplot(121)
 149     else:
 150         ax = fig.add_subplot(111)
 151     ax.set_aspect("equal")
 152     ax.set_title("Final path")
 153     ax.set_xlabel("x [m]")
 154     ax.set_ylabel("y [m]")
 155
 156     for o in s["obst"]:
 157         try:
 158             plt.plot(*plot_nodes(o["segment"]), color="black")
 159         except:
 160             pass
 161         try:
 162             ax.add_artist(plt.Circle((o["circle"][0], o["circle"][1]),
 163                     o["circle"][2],
 164                     color="black", fill=False, hatch="//"))
 165         except:
 166             pass
 167     if PLOT["node"]:
 168         try:
 169             plt.plot(*plot_nodes(t["node"]), color=COLOR["node"],
 170                     marker=".", linestyle = "None")
 171         except:
 172             print("No RRTNode")
 173     if False:
 174         try:
 175             for edges in t["edge"]:
 176                 for e in edges:
 177                     plt.plot([e[0][0], e[1][0]], [e[0][1], e[1][1]],
 178                             color=COLOR["edge"])
 179         except:
 180             print("No edges")
 181     if PLOT["sample"]:
 182         try:
 183             if PLOT["frame"]:
 184                 for i in t["samp"]:
 185                     plt.plot(*car_frame(i), color=COLOR["sample"])
 186             else:
 187                 plt.plot(*plot_nodes(t["samp"]), color=COLOR["sample"],
 188                         marker=".", linestyle = "None")
 189         except:
 190             print("No RRTSample")
 191     if PLOT["path"]:
 192         try:
 193             for path in range(len(t["path"])):
 194                 plt.plot(*plot_nodes(t["path"][path]), color=COLOR["path"])
 195         except:
 196             print("No path")
 197     if PLOT["traj"]:
 198         try:
 199             for traj in range(len(t["traj"])):
 200                 if PLOT["frame"]:
 201                     for i in t["traj"][traj]:
 202                         plt.plot(*car_frame(i), color=COLOR["log"][traj],
 203                                 label=t["cost"][traj], lw=1)
 204                 else:
 205                     try:
 206                         plt.plot(
 207                                 *plot_nodes(t["traj"][traj]),
 208                                 color=COLOR["log"][traj],
 209                                 label=t["cost"][traj])
 210                     except:
 211                         plt.plot(
 212                                 *plot_nodes(t["traj"][traj]),
 213                                 color="black",
 214                                 label=t["cost"][traj])
 215         except:
 216             print("No trajectory")
 217
 218     if PLOT["slot"]:
 219         try:
 220             plt.plot(*plot_nodes(s["slot"]["polygon"]), color=COLOR["slot"])
 221         except:
 222             print("No slot")
 223
 224     try: # init
 225         plt.plot(*car_frame(t["init"]), color="red", lw=2)
 226         plt.plot(*plot_nodes([t["init"]]), color="red", marker="+", ms=12)
 227     except:
 228         plt.plot(*car_frame(s["init"]), color="red", lw=2)
 229         plt.plot(*plot_nodes([s["init"]]), color="red", marker="+", ms=12)
 230     try: # goal
 231         plt.plot(*plot_nodes([t["goal"]]), color="red", marker="+", ms=12)
 232         plt.plot(*car_frame(t["goal"]), color="red", lw=2)
 233     except:
 234         plt.plot(*plot_nodes([s["goal"]]), color="red", marker="+", ms=12)
 235         plt.plot(*car_frame(s["goal"]), color="red", lw=2)
 236     try: # middle
 237         plt.plot(*car_frame(t["midd"]), color="red", lw=2)
 238         plt.plot(*plot_nodes([t["midd"]]), color="red", marker="+", ms=12)
 239     except:
 240         pass
 241
 242     #plt.fill_between(
 243     #            [-2.7, 3.25, 3.25, -2.7],
 244     #            [-0.5, -0.5, 0, 0],
 245     #            [26, 26, 26.5, 26.5],
 246     #            facecolor="none", hatch="//", edgecolor="black", linewidth=0)
 247     #plt.fill_between(
 248     #            [-2.7, -2.7, -0, -0, -2.2, -2.2, 0, 0, -2.7],
 249     #            [-0.5, 26.5, 26.5, 19.5, 19.5, 13, 13, -0.5, -0.5],
 250     #            facecolor="none", hatch="//", edgecolor="black", linewidth=0)
 251     #plt.fill_between(
 252     #            [3.25, 3.25, 2.75, 2.75, 3.25],
 253     #            [-0.5, 26.5, 26.5, -0.5, -0.5],
 254     #            facecolor="none", hatch="//", edgecolor="black", linewidth=0)
 255
 256     #plt.text(1, 0.2, s="1", color="red")
 257     #plt.text(2, 12.2, s="2", color="red")
 258     #plt.text(3.6, 10, s="3", color="black")
 259     #plt.text(-0.5, 14.3, s="4", color="black")
 260     #plt.text(3, 18, s="5", color="blue")
 261     #plt.text(-4.7, 0.8, s="6", color="orange")
 262     #plt.text(-2, 11, s="7", color="gray")
 263
 264     handles, labels = ax.get_legend_handles_labels()
 265
 266     if not "edge" in t.keys():
 267         plt.show()
 268         plt.close(fig)
 269         exit(0)
 270
 271     ################
 272     ## 2nd subplot
 273
 274     ax = fig.add_subplot(122)
 275     ax.set_aspect("equal")
 276     ax.set_title("All edges")
 277     ax.set_xlabel("x [m]")
 278     ax.set_ylabel("y [m]")
 279
 280     for o in s["obst"]:
 281         try:
 282             plt.plot(*plot_nodes(o["segment"]), color="black")
 283         except:
 284             pass
 285         try:
 286             ax.add_artist(plt.Circle((o["circle"][0], o["circle"][1]),
 287                     o["circle"][2],
 288                     color="black", fill=False, hatch="//"))
 289         except:
 290             pass
 291     if PLOT["node"]:
 292         try:
 293             plt.plot(*plot_nodes(t["node"]), color=COLOR["node"],
 294                     marker=".", linestyle = "None")
 295         except:
 296             print("No RRTNode")
 297     if True:
 298         try:
 299             for edges in t["edge"]:
 300                 for e in edges:
 301                     plt.plot([e[0][0], e[1][0]], [e[0][1], e[1][1]],
 302                             color=COLOR["edge"], lw=1)
 303         except:
 304             print("No edges")
 305     if PLOT["sample"]:
 306         try:
 307             if PLOT["frame"]:
 308                 for i in t["samp"]:
 309                     plt.plot(*car_frame(i), color=COLOR["sample"])
 310             else:
 311                 plt.plot(*plot_nodes(t["samp"]), color=COLOR["sample"],
 312                         marker=".", linestyle = "None")
 313         except:
 314             print("No RRTSample")
 315     if PLOT["path"]:
 316         try:
 317             for path in range(len(t["path"])):
 318                 plt.plot(*plot_nodes(t["path"][path]), color=COLOR["path"])
 319         except:
 320             print("No path")
 321     if False:
 322         try:
 323             for traj in range(len(t["traj"])):
 324                 if PLOT["frame"]:
 325                     for i in t["traj"][traj]:
 326                         plt.plot(*car_frame(i), color=COLOR["log"][traj],
 327                                 label=t["cost"][traj])
 328                 else:
 329                     try:
 330                         plt.plot(
 331                                 *plot_nodes(t["traj"][traj]),
 332                                 color=COLOR["log"][traj],
 333                                 label=t["cost"][traj])
 334                     except:
 335                         plt.plot(
 336                                 *plot_nodes(t["traj"][traj]),
 337                                 color="black",
 338                                 label=t["cost"][traj])
 339         except:
 340             print("No trajectory")
 341
 342     if PLOT["slot"]:
 343         try:
 344             plt.plot(*plot_nodes(s["slot"]["polygon"]), color=COLOR["slot"])
 345         except:
 346             print("No slot")
 347
 348     plt.plot(*car_frame(s["init"]), color="red", lw=2)
 349     plt.plot(*car_frame(s["goal"]), color="red", lw=2)
 350     plt.plot(*plot_nodes([s["init"]]), color="red", marker="+", ms=12)
 351     plt.plot(*plot_nodes([s["goal"]]), color="red", marker="+", ms=12)
 352
 353     #plt.fill_between(
 354     #            [-2.7, 3.25, 3.25, -2.7],
 355     #            [-0.5, -0.5, 0, 0],
 356     #            [26, 26, 26.5, 26.5],
 357     #            facecolor="none", hatch="//", edgecolor="black", linewidth=0)
 358     #plt.fill_between(
 359     #            [-2.7, -2.7, -0, -0, -2.2, -2.2, 0, 0, -2.7],
 360     #            [-0.5, 26.5, 26.5, 19.5, 19.5, 13, 13, -0.5, -0.5],
 361     #            facecolor="none", hatch="//", edgecolor="black", linewidth=0)
 362     #plt.fill_between(
 363     #            [3.25, 3.25, 2.75, 2.75, 3.25],
 364     #            [-0.5, 26.5, 26.5, -0.5, -0.5],
 365     #            facecolor="none", hatch="//", edgecolor="black", linewidth=0)
 366
 367     handles, labels = ax.get_legend_handles_labels()
 368
 369     # END OF SUBPLOTS
 370     plt.show()
 371     plt.close(fig)