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
  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     ax = fig.add_subplot(121)
 146     ax.set_aspect("equal")
 147     ax.set_title("Final path")
 148     ax.set_xlabel("x [m]")
 149     ax.set_ylabel("y [m]")
 150
 151     for o in s["obst"]:
 152         try:
 153             plt.plot(*plot_nodes(o["segment"]), color="black")
 154         except:
 155             pass
 156         try:
 157             ax.add_artist(plt.Circle((o["circle"][0], o["circle"][1]),
 158                     o["circle"][2],
 159                     color="black", fill=False, hatch="//"))
 160         except:
 161             pass
 162     if PLOT["node"]:
 163         try:
 164             plt.plot(*plot_nodes(t["node"]), color=COLOR["node"],
 165                     marker=".", linestyle = "None")
 166         except:
 167             print("No RRTNode")
 168     if False:
 169         try:
 170             for edges in t["edge"]:
 171                 for e in edges:
 172                     plt.plot([e[0][0], e[1][0]], [e[0][1], e[1][1]],
 173                             color=COLOR["edge"])
 174         except:
 175             print("No edges")
 176     if PLOT["sample"]:
 177         try:
 178             if PLOT["frame"]:
 179                 for i in t["samp"]:
 180                     plt.plot(*car_frame(i), color=COLOR["sample"])
 181             else:
 182                 plt.plot(*plot_nodes(t["samp"]), color=COLOR["sample"],
 183                         marker=".", linestyle = "None")
 184         except:
 185             print("No RRTSample")
 186     if PLOT["path"]:
 187         try:
 188             for path in range(len(t["path"])):
 189                 plt.plot(*plot_nodes(t["path"][path]), color=COLOR["path"])
 190         except:
 191             print("No path")
 192     if PLOT["traj"]:
 193         try:
 194             for traj in range(len(t["traj"])):
 195                 if PLOT["frame"]:
 196                     for i in t["traj"][traj]:
 197                         plt.plot(*car_frame(i), color=COLOR["log"][traj],
 198                                 label=t["cost"][traj], lw=1)
 199                 else:
 200                     try:
 201                         plt.plot(
 202                                 *plot_nodes(t["traj"][traj]),
 203                                 color=COLOR["log"][traj],
 204                                 label=t["cost"][traj])
 205                     except:
 206                         plt.plot(
 207                                 *plot_nodes(t["traj"][traj]),
 208                                 color="black",
 209                                 label=t["cost"][traj])
 210         except:
 211             print("No trajectory")
 212
 213     plt.plot(*car_frame(s["init"]), color="red", lw=2)
 214     plt.plot(*car_frame(s["goal"]), color="red", lw=2)
 215     plt.plot(*plot_nodes([s["init"]]), color="red", marker="+", ms=12)
 216     plt.plot(*plot_nodes([s["goal"]]), color="red", marker="+", ms=12)
 217
 218     plt.fill_between(
 219                 [-2.7, 3.25, 3.25, -2.7],
 220                 [1.5, 1.5, 2, 2],
 221                 [26, 26, 26.5, 26.5],
 222                 facecolor="none", hatch="//", edgecolor="black", linewidth=0)
 223     plt.fill_between(
 224                 [-2.7, -2.7, -0, -0, -2.2, -2.2, 0, 0, -2.7],
 225                 [1.5, 26.5, 26.5, 19.5, 19.5, 13, 13, 1.5, 1.5],
 226                 facecolor="none", hatch="//", edgecolor="black", linewidth=0)
 227     plt.fill_between(
 228                 [3.25, 3.25, 2.75, 2.75, 3.25],
 229                 [1.5, 26.5, 26.5, 1.5, 1.5],
 230                 facecolor="none", hatch="//", edgecolor="black", linewidth=0)
 231
 232     #plt.text(1, 0.2, s="1", color="red")
 233     #plt.text(2, 12.2, s="2", color="red")
 234     #plt.text(3.6, 10, s="3", color="black")
 235     #plt.text(-0.5, 14.3, s="4", color="black")
 236     #plt.text(3, 18, s="5", color="blue")
 237     #plt.text(-4.7, 0.8, s="6", color="orange")
 238     #plt.text(-2, 11, s="7", color="gray")
 239
 240     handles, labels = ax.get_legend_handles_labels()
 241
 242     ################
 243     ## 2nd subplot
 244
 245     ax = fig.add_subplot(122)
 246     ax.set_aspect("equal")
 247     ax.set_title("All edges")
 248     ax.set_xlabel("x [m]")
 249     ax.set_ylabel("y [m]")
 250
 251     for o in s["obst"]:
 252         try:
 253             plt.plot(*plot_nodes(o["segment"]), color="black")
 254         except:
 255             pass
 256         try:
 257             ax.add_artist(plt.Circle((o["circle"][0], o["circle"][1]),
 258                     o["circle"][2],
 259                     color="black", fill=False, hatch="//"))
 260         except:
 261             pass
 262     if PLOT["node"]:
 263         try:
 264             plt.plot(*plot_nodes(t["node"]), color=COLOR["node"],
 265                     marker=".", linestyle = "None")
 266         except:
 267             print("No RRTNode")
 268     if True:
 269         try:
 270             for edges in t["edge"]:
 271                 for e in edges:
 272                     plt.plot([e[0][0], e[1][0]], [e[0][1], e[1][1]],
 273                             color=COLOR["edge"], lw=1)
 274         except:
 275             print("No edges")
 276     if PLOT["sample"]:
 277         try:
 278             if PLOT["frame"]:
 279                 for i in t["samp"]:
 280                     plt.plot(*car_frame(i), color=COLOR["sample"])
 281             else:
 282                 plt.plot(*plot_nodes(t["samp"]), color=COLOR["sample"],
 283                         marker=".", linestyle = "None")
 284         except:
 285             print("No RRTSample")
 286     if PLOT["path"]:
 287         try:
 288             for path in range(len(t["path"])):
 289                 plt.plot(*plot_nodes(t["path"][path]), color=COLOR["path"])
 290         except:
 291             print("No path")
 292     if False:
 293         try:
 294             for traj in range(len(t["traj"])):
 295                 if PLOT["frame"]:
 296                     for i in t["traj"][traj]:
 297                         plt.plot(*car_frame(i), color=COLOR["log"][traj],
 298                                 label=t["cost"][traj])
 299                 else:
 300                     try:
 301                         plt.plot(
 302                                 *plot_nodes(t["traj"][traj]),
 303                                 color=COLOR["log"][traj],
 304                                 label=t["cost"][traj])
 305                     except:
 306                         plt.plot(
 307                                 *plot_nodes(t["traj"][traj]),
 308                                 color="black",
 309                                 label=t["cost"][traj])
 310         except:
 311             print("No trajectory")
 312
 313     plt.plot(*car_frame(s["init"]), color="red", lw=2)
 314     plt.plot(*car_frame(s["goal"]), color="red", lw=2)
 315     plt.plot(*plot_nodes([s["init"]]), color="red", marker="+", ms=12)
 316     plt.plot(*plot_nodes([s["goal"]]), color="red", marker="+", ms=12)
 317
 318     plt.fill_between(
 319                 [-2.7, 3.25, 3.25, -2.7],
 320                 [1.5, 1.5, 2, 2],
 321                 [26, 26, 26.5, 26.5],
 322                 facecolor="none", hatch="//", edgecolor="black", linewidth=0)
 323     plt.fill_between(
 324                 [-2.7, -2.7, -0, -0, -2.2, -2.2, 0, 0, -2.7],
 325                 [1.5, 26.5, 26.5, 19.5, 19.5, 13, 13, 1.5, 1.5],
 326                 facecolor="none", hatch="//", edgecolor="black", linewidth=0)
 327     plt.fill_between(
 328                 [3.25, 3.25, 2.75, 2.75, 3.25],
 329                 [1.5, 26.5, 26.5, 1.5, 1.5],
 330                 facecolor="none", hatch="//", edgecolor="black", linewidth=0)
 331
 332     handles, labels = ax.get_legend_handles_labels()
 333
 334     # END OF SUBPLOTS
 335     plt.show()
 336     plt.close(fig)