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.418
  11 LENGTH = 4.970
  12 SAFETY_DIST = 0
  13 WHEEL_BASE = 2.9
  14 WIDTH = 1.931
  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": "lightgrey",
  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     lfx = pose[0]
  47     lfx += (WIDTH / 2.0) * cos(pose[2] + pi / 2.0)
  48     lfx += WHEEL_BASE * cos(pose[2])
  49     lfx += SAFETY_DIST * cos(pose[2])
  50
  51     lrx = pose[0]
  52     lrx += (WIDTH / 2.0) * cos(pose[2] + pi / 2.0)
  53     lrx += -SAFETY_DIST * cos(pose[2])
  54
  55     rrx = pose[0]
  56     rrx += (WIDTH / 2.0) * cos(pose[2] - pi / 2.0)
  57     rrx += -SAFETY_DIST * cos(pose[2])
  58
  59     rfx = pose[0]
  60     rfx += (WIDTH / 2.0) * cos(pose[2] - pi / 2.0)
  61     rfx += WHEEL_BASE * cos(pose[2])
  62     rfx += SAFETY_DIST * cos(pose[2])
  63
  64     lfy = pose[1]
  65     lfy += (WIDTH / 2.0) * sin(pose[2] + pi / 2.0)
  66     lfy += WHEEL_BASE * sin(pose[2])
  67     lfy += SAFETY_DIST * sin(pose[2])
  68
  69     lry = pose[1]
  70     lry += (WIDTH / 2.0) * sin(pose[2] + pi / 2.0)
  71     lry += -SAFETY_DIST * sin(pose[2])
  72
  73     rry = pose[1]
  74     rry += (WIDTH / 2.0) * sin(pose[2] - pi / 2.0)
  75     rry += -SAFETY_DIST * sin(pose[2])
  76
  77     rfy = pose[1]
  78     rfy += (WIDTH / 2.0) * sin(pose[2] - pi / 2.0)
  79     rfy += WHEEL_BASE * sin(pose[2])
  80     rfy += SAFETY_DIST * sin(pose[2])
  81
  82     xcoords = (lfx, lrx, rrx, rfx)
  83     ycoords = (lfy, lry, rry, rfy)
  84     return (xcoords, ycoords)
  85
  86 def load_scenario(fname):
  87     """Load scenario from file."""
  88     if fname is None:
  89         raise ValueError("File name as argument needed")
  90     with open(fname, "r") as f:
  91         scenario = loads(f.read())
  92     return scenario
  93
  94 def load_trajectory(fname):
  95     """Load trajectory from file."""
  96     if fname is None:
  97         raise ValueError("File name as argument needed")
  98     with open(fname, "r") as f:
  99         trajectory = loads(f.read())
 100         return trajectory
 101
 102 def plot_nodes(nodes=[]):
 103     """Return xcoords, ycoords of nodes to plot.
 104
 105     Keyword arguments:
 106     nodes -- The list of nodes to plot.
 107     """
 108     xcoords = []
 109     ycoords = []
 110     for n in nodes:
 111         xcoords.append(n[0])
 112         ycoords.append(n[1])
 113     return (xcoords, ycoords)
 114
 115 def plot_segments(segments=[]):
 116     """Return xcoords, ycoords of segments to plot.
 117
 118     Keyword arguments:
 119     segments -- The list of segments to plot.
 120     """
 121     pass
 122
 123 if __name__ == "__main__":
 124     s = load_scenario(SCEN_FILE)
 125     try:
 126         t = load_trajectory(TRAJ_FILE) # fixed to trajectories
 127     except:
 128         pass
 129
 130     plt.rcParams["font.size"] = 24
 131     fig = plt.figure()
 132     ax = fig.add_subplot(111)
 133     ax.set_aspect("equal")
 134     ax.set_title("SCENARIO")
 135     ax.set_xlabel("Longitudinal direction [m]")
 136     ax.set_ylabel("Lateral direction [m]")
 137
 138     # plot here
 139     for o in s["obst"]:
 140         try:
 141             plt.plot(*plot_nodes(o["segment"]), color="black")
 142         except:
 143             pass
 144         try:
 145             ax.add_artist(plt.Circle((o["circle"][0], o["circle"][1]),
 146                     o["circle"][2],
 147                     color="black", fill=False))
 148         except:
 149             pass
 150     if PLOT["node"]:
 151         try:
 152             plt.plot(*plot_nodes(t["node"]), color=COLOR["node"],
 153                     marker=".", linestyle = "None")
 154         except:
 155             print("No RRTNode")
 156     if PLOT["edge"]:
 157         try:
 158             for edges in t["edge"]:
 159                 for e in edges:
 160                     plt.plot([e[0][0], e[1][0]], [e[0][1], e[1][1]],
 161                             color=COLOR["edge"])
 162         except:
 163             print("No edges")
 164     if PLOT["sample"]:
 165         try:
 166             if PLOT["frame"]:
 167                 for i in t["samp"]:
 168                     plt.plot(*car_frame(i), color=COLOR["sample"])
 169             else:
 170                 plt.plot(*plot_nodes(t["samp"]), color=COLOR["sample"],
 171                         marker=".", linestyle = "None")
 172         except:
 173             print("No RRTSample")
 174     if PLOT["path"]:
 175         try:
 176             for path in range(len(t["path"])):
 177                 plt.plot(*plot_nodes(t["path"][path]), color=COLOR["path"])
 178         except:
 179             print("No path")
 180     if PLOT["traj"]:
 181         try:
 182             for traj in range(len(t["traj"])):
 183                 if PLOT["frame"]:
 184                     for i in t["traj"][traj]:
 185                         plt.plot(*car_frame(i), color=COLOR["log"][traj],
 186                                 label=t["cost"][traj])
 187                 else:
 188                     try:
 189                         plt.plot(
 190                                 *plot_nodes(t["traj"][traj]),
 191                                 color=COLOR["log"][traj],
 192                                 label=t["cost"][traj])
 193                     except:
 194                         plt.plot(
 195                                 *plot_nodes(t["traj"][traj]),
 196                                 color="black",
 197                                 label=t["cost"][traj])
 198         except:
 199             print("No trajectory")
 200         for i in range(len(t["traj"][-1]) - 1):
 201             n1 = t["traj"][-1][i]
 202             n2 = t["traj"][-1][i + 1]
 203             if (sign(n1[4]) != sign(n2[4])):
 204                 plt.plot(n2[0], n2[1], color=COLOR["log"][1], marker=".")
 205     plt.plot(*plot_nodes([s["init"]]), color=COLOR["start"], marker=".")
 206     plt.plot(*plot_nodes([s["goal"]]), color=COLOR["goal"], marker=".")
 207     # end plot here
 208
 209     handles, labels = ax.get_legend_handles_labels()
 210     lgd = ax.legend(handles, labels, loc="upper center",
 211             bbox_to_anchor=(0.5, -0.11), title="Cost")
 212     plt.show()
 213     #plt.savefig("{}.png".format(argv[2]), bbox_extra_artists=(lgd,),
 214     #        bbox_inches='tight')
 215     plt.close(fig)