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