first release

PathPlanning/ProbabilisticRoadMap/probabilistic_road_map.py

@@ -9,18 +9,21 @@ author: Atsushi Sakai (@Atsushi_twi)
 import random
 import math
 import numpy as np
+import scipy.spatial
 import matplotlib.pyplot as plt
-from matplotrecorder import matplotrecorder
-from pyfastnns import pyfastnns
-matplotrecorder.donothing = True
 
 # parameter
-N_SAMPLE = 500
-N_KNN = 10
+N_SAMPLE = 500  # number of sample_points
+N_KNN = 10  # number of edge from one sampled point
 MAX_EDGE_LEN = 30.0  # [m] Maximum edge length
 
+show_animation = True
+
 
 class Node:
+    """
+    Node class for dijkstra search
+    """
 
     def __init__(self, x, y, cost, pind):
         self.x = x
@@ -32,12 +35,55 @@ class Node:
         return str(self.x) + "," + str(self.y) + "," + str(self.cost) + "," + str(self.pind)
 
 
+class KDTree:
+    """
+    Nearest neighbor search class with KDTree
+    """
+
+    def __init__(self, data):
+        # store kd-tree
+        self.tree = scipy.spatial.cKDTree(data)
+
+    def search(self, inp, k=1):
+        u"""
+        Search NN
+
+        inp: input data, single frame or multi frame
+
+        """
+
+        if len(inp.shape) >= 2:  # multi input
+            index = []
+            dist = []
+
+            for i in inp.T:
+                idist, iindex = self.tree.query(i, k=k)
+                index.append(iindex)
+                dist.append(idist)
+
+            return index, dist
+        else:
+            dist, index = self.tree.query(inp, k=k)
+            return index, dist
+
+    def search_in_distance(self, inp, r):
+        u"""
+        find points with in a distance r
+        """
+
+        index = self.tree.query_ball_point(inp, r)
+        return index
+
+
 def PRM_planning(sx, sy, gx, gy, ox, oy, rr):
 
-    sample_x, sample_y = sample_points(sx, sy, gx, gy, rr, ox, oy)
-    plt.plot(sample_x, sample_y, ".r")
+    obkdtree = KDTree(np.vstack((ox, oy)).T)
 
-    road_map = generate_roadmap(sample_x, sample_y, rr)
+    sample_x, sample_y = sample_points(sx, sy, gx, gy, rr, ox, oy, obkdtree)
+    if show_animation:
+        plt.plot(sample_x, sample_y, ".b")
+
+    road_map = generate_roadmap(sample_x, sample_y, rr, obkdtree)
 
     rx, ry = dijkstra_planning(
         sx, sy, gx, gy, ox, oy, rr, road_map, sample_x, sample_y)
@@ -45,24 +91,64 @@ def PRM_planning(sx, sy, gx, gy, ox, oy, rr):
     return rx, ry
 
 
-def generate_roadmap(sample_x, sample_y, rr):
+def is_collision(sx, sy, gx, gy, rr, okdtree):
+    x = sx
+    y = sy
+    dx = gx - sx
+    dy = gy - sy
+    yaw = math.atan2(gy - sy, gx - sx)
+    d = math.sqrt(dx**2 + dy**2)
+
+    if d >= MAX_EDGE_LEN:
+        return True
+
+    D = rr
+    nstep = round(d / D)
+
+    for i in range(nstep):
+        idxs, dist = okdtree.search(np.matrix([x, y]).T)
+        if dist[0] <= rr:
+            return True  # collision
+        x += D * math.cos(yaw)
+        y += D * math.sin(yaw)
+
+    # goal point check
+    idxs, dist = okdtree.search(np.matrix([gx, gy]).T)
+    if dist[0] <= rr:
+        return True  # collision
+
+    return False  # OK
+
+
+def generate_roadmap(sample_x, sample_y, rr, obkdtree):
+    """
+    Road map generation
+
+    sample_x: [m] x positions of sampled points
+    sample_y: [m] y positions of sampled points
+    rr: Robot Radius[m]
+    obkdtree: KDTree object of obstacles
+    """
 
     road_map = []
     nsample = len(sample_x)
-    skdtree = pyfastnns.NNS(np.vstack((sample_x, sample_y)).T)
+    skdtree = KDTree(np.vstack((sample_x, sample_y)).T)
 
     for (i, ix, iy) in zip(range(nsample), sample_x, sample_y):
 
-        index = skdtree.search(
+        index, dists = skdtree.search(
             np.matrix([ix, iy]).T, k=nsample)
+        inds = index[0][0]
         edge_id = []
+        #  print(index)
 
-        for ii in range(1, len(index[0][0][0])):
-            #  nx = sample_x[index[i]]
-            #  ny = sample_y[index[i]]
+        for ii in range(1, len(inds)):
+            nx = sample_x[inds[ii]]
+            ny = sample_y[inds[ii]]
+
+            if not is_collision(ix, iy, nx, ny, rr, obkdtree):
+                edge_id.append(inds[ii])
 
-            #  if !is_collision(ix, iy, nx, ny, rr, okdtree)
-            edge_id.append(index[0][0][0][ii])
             if len(edge_id) >= N_KNN:
                 break
 
@@ -94,21 +180,16 @@ def dijkstra_planning(sx, sy, gx, gy, ox, oy, rr, road_map, sample_x, sample_y):
             print("Cannot find path")
             break
 
-        print(len(openset), len(closedset))
-
         c_id = min(openset, key=lambda o: openset[o].cost)
         current = openset[c_id]
-        print("current", current, c_id)
-        #  input()
 
         # show graph
-        plt.plot(current.x, current.y, "xc")
-        if len(closedset.keys()) % 10 == 0:
+        if show_animation and len(closedset.keys()) % 2 == 0:
+            plt.plot(current.x, current.y, "xg")
             plt.pause(0.001)
-            matplotrecorder.save_frame()
 
         if c_id == (len(road_map) - 1):
-            print("Find goal")
+            print("goal is found!")
             ngoal.pind = current.pind
             ngoal.cost = current.cost
             break
@@ -121,16 +202,12 @@ def dijkstra_planning(sx, sy, gx, gy, ox, oy, rr, road_map, sample_x, sample_y):
         # expand search grid based on motion model
         for i in range(len(road_map[c_id])):
             n_id = road_map[c_id][i]
-            print(i, n_id)
             dx = sample_x[n_id] - current.x
             dy = sample_y[n_id] - current.y
             d = math.sqrt(dx**2 + dy**2)
             node = Node(sample_x[n_id], sample_y[n_id],
                         current.cost + d, c_id)
 
-            #  if not verify_node(node, obmap, minx, miny, maxx, maxy):
-            #  continue
-
             if n_id in closedset:
                 continue
             # Otherwise if it is already in the open set
@@ -163,7 +240,7 @@ def plot_road_map(road_map, sample_x, sample_y):
                      [sample_y[i], sample_y[ind]], "-k")
 
 
-def sample_points(sx, sy, gx, gy, rr, ox, oy):
+def sample_points(sx, sy, gx, gy, rr, ox, oy, obkdtree):
     maxx = max(ox)
     maxy = max(oy)
     minx = min(ox)
@@ -171,13 +248,11 @@ def sample_points(sx, sy, gx, gy, rr, ox, oy):
 
     sample_x, sample_y = [], []
 
-    nns = pyfastnns.NNS(np.vstack((ox, oy)).T)
-
     while len(sample_x) <= N_SAMPLE:
         tx = (random.random() - minx) * (maxx - minx)
         ty = (random.random() - miny) * (maxy - miny)
 
-        index, dist = nns.search(np.matrix([tx, ty]).T)
+        index, dist = obkdtree.search(np.matrix([tx, ty]).T)
 
         if dist[0] >= rr:
             sample_x.append(tx)
@@ -224,8 +299,8 @@ def main():
         oy.append(60.0 - i)
 
     plt.plot(ox, oy, ".k")
-    plt.plot(sx, sy, "xr")
-    plt.plot(gx, gy, "xb")
+    plt.plot(sx, sy, "^r")
+    plt.plot(gx, gy, "^c")
     plt.grid(True)
     plt.axis("equal")
 
@@ -233,11 +308,10 @@ def main():
 
     plt.plot(rx, ry, "-r")
 
-    for i in range(20):
-        matplotrecorder.save_frame()
-    plt.show()
+    assert len(rx) != 0, 'Cannot found path'
 
-    matplotrecorder.save_movie("animation.gif", 0.1)
+    if show_animation:
+        plt.show()
 
 
 if __name__ == '__main__':