Plot 2 subplots

[hubacji1/iamcar.git] / plot.py

diff --git a/plot.py b/plot.py
index 63348b7d747b94ca34b3dab3431131aafd343894..c02d9249ec37f97da5d4d5f03b1b61548015d557 100644 (file)
--- a/plot.py
+++ b/plot.py
@@ -30,7 +30,7 @@ COLOR = {
         "start": "violet",
         "goal": "red",
         "sample": "magenta",
-        "path": "lightgrey",
+        "path": "grey",
         "trajectory": "blue",
         "trajectory-frame": "lightblue",
         "obstacle": "black",
@@ -129,9 +129,10 @@ if __name__ == "__main__":
 
     plt.rcParams["font.size"] = 24
     fig = plt.figure()
-    ax = fig.add_subplot(111)
+    ## 1st subplot
+    ax = fig.add_subplot(121)
     ax.set_aspect("equal")
-    ax.set_title("SCENARIO")
+    ax.set_title("RRT* final path")
     ax.set_xlabel("Longitudinal direction [m]")
     ax.set_ylabel("Lateral direction [m]")
 
@@ -153,7 +154,7 @@ if __name__ == "__main__":
                     marker=".", linestyle = "None")
         except:
             print("No RRTNode")
-    if PLOT["edge"]:
+    if False:
         try:
             for edges in t["edge"]:
                 for e in edges:
@@ -178,6 +179,82 @@ if __name__ == "__main__":
         except:
             print("No path")
     if PLOT["traj"]:
+        try:
+            for traj in range(4):#len(t["traj"])):
+                if PLOT["frame"]:
+                    for i in t["traj"][traj]:
+                        plt.plot(*car_frame(i), color=COLOR["log"][traj],
+                                label=t["cost"][traj])
+                else:
+                    try:
+                        plt.plot(
+                                *plot_nodes(t["traj"][traj]),
+                                color=COLOR["log"][traj],
+                                label=t["cost"][traj])
+                    except:
+                        plt.plot(
+                                *plot_nodes(t["traj"][traj]),
+                                color="black",
+                                label=t["cost"][traj])
+        except:
+            print("No trajectory")
+    plt.plot(*plot_nodes([s["init"]]), color=COLOR["start"], marker=".")
+    plt.plot(*plot_nodes([s["goal"]]), color=COLOR["goal"], marker=".")
+    # end plot here
+
+    handles, labels = ax.get_legend_handles_labels()
+    #lgd = ax.legend(handles, labels, loc="upper center",
+    #        bbox_to_anchor=(0.5, -0.11), title="Cost")
+    ## 2nd subplot
+    ax = fig.add_subplot(122)
+    ax.set_aspect("equal")
+    ax.set_title("RRT* all edges")
+    ax.set_xlabel("Longitudinal direction [m]")
+    ax.set_ylabel("Lateral direction [m]")
+
+    # plot here
+    for o in s["obst"]:
+        try:
+            plt.plot(*plot_nodes(o["segment"]), color="black")
+        except:
+            pass
+        try:
+            ax.add_artist(plt.Circle((o["circle"][0], o["circle"][1]),
+                    o["circle"][2],
+                    color="black", fill=False))
+        except:
+            pass
+    if PLOT["node"]:
+        try:
+            plt.plot(*plot_nodes(t["node"]), color=COLOR["node"],
+                    marker=".", linestyle = "None")
+        except:
+            print("No RRTNode")
+    if True:
+        try:
+            for edges in t["edge"]:
+                for e in edges:
+                    plt.plot([e[0][0], e[1][0]], [e[0][1], e[1][1]],
+                            color=COLOR["edge"])
+        except:
+            print("No edges")
+    if PLOT["sample"]:
+        try:
+            if PLOT["frame"]:
+                for i in t["samp"]:
+                    plt.plot(*car_frame(i), color=COLOR["sample"])
+            else:
+                plt.plot(*plot_nodes(t["samp"]), color=COLOR["sample"],
+                        marker=".", linestyle = "None")
+        except:
+            print("No RRTSample")
+    if PLOT["path"]:
+        try:
+            for path in range(len(t["path"])):
+                plt.plot(*plot_nodes(t["path"][path]), color=COLOR["path"])
+        except:
+            print("No path")
+    if False:
         try:
             for traj in range(len(t["traj"])):
                 if PLOT["frame"]:
@@ -202,8 +279,10 @@ if __name__ == "__main__":
     # end plot here
 
     handles, labels = ax.get_legend_handles_labels()
-    lgd = ax.legend(handles, labels, loc="upper center",
-            bbox_to_anchor=(0.5, -0.11), title="Cost")
+    #lgd = ax.legend(handles, labels, loc="upper center",
+    #        bbox_to_anchor=(0.5, -0.11), title="Cost")
+
+    # END OF SUBPLOTS
     plt.show()
     #plt.savefig("{}.png".format(argv[2]), bbox_extra_artists=(lgd,),
     #        bbox_inches='tight')