use scipy kd-tree directly (AtsushiSakai#337)

Author: AtsushiSakai

PathPlanning/HybridAStar/car.py

@@ -31,7 +31,7 @@ def check_car_collision(x_list, y_list, yaw_list, ox, oy, kd_tree):
         cx = i_x + W_BUBBLE_DIST * cos(i_yaw)
         cy = i_y + W_BUBBLE_DIST * sin(i_yaw)
 
-        ids = kd_tree.search_in_distance([cx, cy], W_BUBBLE_R)
+        ids = kd_tree.query_ball_point([cx, cy], W_BUBBLE_R)
 
         if not ids:
             continue
@@ -71,7 +71,7 @@ def plot_arrow(x, y, yaw, length=1.0, width=0.5, fc="r", ec="k"):
 def plot_car(x, y, yaw):
     car_color = '-k'
     c, s = cos(yaw), sin(yaw)
-    rot = Rot.from_euler('z', yaw).as_matrix()[0:2, 0:2]
+    rot = Rot.from_euler('z', -yaw).as_matrix()[0:2, 0:2]
     car_outline_x, car_outline_y = [], []
     for rx, ry in zip(VRX, VRY):
         converted_xy = np.stack([rx, ry]).T @ rot

PathPlanning/HybridAStar/hybrid_a_star.py

@@ -13,7 +13,7 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
-import scipy.spatial
+from scipy.spatial import cKDTree
 
 sys.path.append(os.path.dirname(os.path.abspath(__file__))
                 + "/../ReedsSheppPath")
@@ -67,44 +67,6 @@ def __init__(self, x_list, y_list, yaw_list, direction_list, cost):
         self.cost = cost
 
 
-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):
-        """
-        Search NN
-        inp: input data, single frame or multi frame
-        """
-
-        if len(inp.shape) >= 2:  # multi input
-            index = []
-            dist = []
-
-            for i in inp.T:
-                i_dist, i_index = self.tree.query(i, k=k)
-                index.append(i_index)
-                dist.append(i_dist)
-
-            return index, dist
-
-        dist, index = self.tree.query(inp, k=k)
-        return index, dist
-
-    def search_in_distance(self, inp, r):
-        """
-        find points with in a distance r
-        """
-
-        index = self.tree.query_ball_point(inp, r)
-        return index
-
-
 class Config:
 
     def __init__(self, ox, oy, xy_resolution, yaw_resolution):
@@ -297,7 +259,7 @@ def hybrid_a_star_planning(start, goal, ox, oy, xy_resolution, yaw_resolution):
     start[2], goal[2] = rs.pi_2_pi(start[2]), rs.pi_2_pi(goal[2])
     tox, toy = ox[:], oy[:]
 
-    obstacle_kd_tree = KDTree(np.vstack((tox, toy)).T)
+    obstacle_kd_tree = cKDTree(np.vstack((tox, toy)).T)
 
     config = Config(tox, toy, xy_resolution, yaw_resolution)
 

PathPlanning/ProbabilisticRoadMap/probabilistic_road_map.py

@@ -9,8 +9,8 @@
 import random
 import math
 import numpy as np
-import scipy.spatial
 import matplotlib.pyplot as plt
+from scipy.spatial import cKDTree
 
 # parameter
 N_SAMPLE = 500  # number of sample_points
@@ -36,49 +36,9 @@ def __str__(self):
                str(self.cost) + "," + str(self.parent_index)
 
 
-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):
-        """
-        Search NN
-
-        inp: input data, single frame or multi frame
-
-        """
-
-        if len(inp.shape) >= 2:  # multi input
-            index = []
-            dist = []
-
-            for i in inp.T:
-                i_dist, i_index = self.tree.query(i, k=k)
-                index.append(i_index)
-                dist.append(i_dist)
-
-            return index, dist
-
-        dist, index = self.tree.query(inp, k=k)
-        return index, dist
-
-    def search_in_distance(self, inp, r):
-        """
-        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):
 
-    obstacle_kd_tree = KDTree(np.vstack((ox, oy)).T)
+    obstacle_kd_tree = cKDTree(np.vstack((ox, oy)).T)
 
     sample_x, sample_y = sample_points(sx, sy, gx, gy,
                                        rr, ox, oy, obstacle_kd_tree)
@@ -108,15 +68,15 @@ def is_collision(sx, sy, gx, gy, rr, obstacle_kd_tree):
     n_step = round(d / D)
 
     for i in range(n_step):
-        _, dist = obstacle_kd_tree.search(np.array([x, y]).reshape(2, 1))
-        if dist[0] <= rr:
+        dist, _ = obstacle_kd_tree.query([x, y])
+        if dist <= rr:
             return True  # collision
         x += D * math.cos(yaw)
         y += D * math.sin(yaw)
 
     # goal point check
-    _, dist = obstacle_kd_tree.search(np.array([gx, gy]).reshape(2, 1))
-    if dist[0] <= rr:
+    dist, _ = obstacle_kd_tree.query([gx, gy])
+    if dist <= rr:
         return True  # collision
 
     return False  # OK
@@ -134,22 +94,19 @@ def generate_road_map(sample_x, sample_y, rr, obstacle_kd_tree):
 
     road_map = []
     n_sample = len(sample_x)
-    sample_kd_tree = KDTree(np.vstack((sample_x, sample_y)).T)
+    sample_kd_tree = cKDTree(np.vstack((sample_x, sample_y)).T)
 
     for (i, ix, iy) in zip(range(n_sample), sample_x, sample_y):
 
-        index, dists = sample_kd_tree.search(
-            np.array([ix, iy]).reshape(2, 1), k=n_sample)
-        inds = index[0]
+        dists, indexes = sample_kd_tree.query([ix, iy], k=n_sample)
         edge_id = []
-        #  print(index)
 
-        for ii in range(1, len(inds)):
-            nx = sample_x[inds[ii]]
-            ny = sample_y[inds[ii]]
+        for ii in range(1, len(indexes)):
+            nx = sample_x[indexes[ii]]
+            ny = sample_y[indexes[ii]]
 
             if not is_collision(ix, iy, nx, ny, rr, obstacle_kd_tree):
-                edge_id.append(inds[ii])
+                edge_id.append(indexes[ii])
 
             if len(edge_id) >= N_KNN:
                 break
@@ -270,9 +227,9 @@ def sample_points(sx, sy, gx, gy, rr, ox, oy, obstacle_kd_tree):
         tx = (random.random() * (max_x - min_x)) + min_x
         ty = (random.random() * (max_y - min_y)) + min_y
 
-        index, dist = obstacle_kd_tree.search(np.array([tx, ty]).reshape(2, 1))
+        dist, index = obstacle_kd_tree.query([tx, ty])
 
-        if dist[0] >= rr:
+        if dist >= rr:
             sample_x.append(tx)
             sample_y.append(ty)
 

PathPlanning/VoronoiRoadMap/kdtree.py

@@ -8,10 +8,9 @@
 
 import math
 import numpy as np
-import scipy.spatial
 import matplotlib.pyplot as plt
 from dijkstra_search import DijkstraSearch
-from kdtree import KDTree
+from scipy.spatial import cKDTree, Voronoi
 
 show_animation = True
 
@@ -24,7 +23,7 @@ def __init__(self):
         self.MAX_EDGE_LEN = 30.0  # [m] Maximum edge length
 
     def planning(self, sx, sy, gx, gy, ox, oy, robot_radius):
-        obstacle_tree = KDTree(np.vstack((ox, oy)).T)
+        obstacle_tree = cKDTree(np.vstack((ox, oy)).T)
 
         sample_x, sample_y = self.voronoi_sampling(sx, sy, gx, gy, ox, oy)
         if show_animation:  # pragma: no cover
@@ -53,15 +52,15 @@ def is_collision(self, sx, sy, gx, gy, rr, obstacle_kd_tree):
         n_step = round(d / D)
 
         for i in range(n_step):
-            ids, dist = obstacle_kd_tree.search(np.array([x, y]).reshape(2, 1))
-            if dist[0] <= rr:
+            dist, _ = obstacle_kd_tree.query([x, y])
+            if dist <= rr:
                 return True  # collision
             x += D * math.cos(yaw)
             y += D * math.sin(yaw)
 
         # goal point check
-        ids, dist = obstacle_kd_tree.search(np.array([gx, gy]).reshape(2, 1))
-        if dist[0] <= rr:
+        dist, _ = obstacle_kd_tree.query([gx, gy])
+        if dist <= rr:
             return True  # collision
 
         return False  # OK
@@ -78,22 +77,20 @@ def generate_road_map_info(self, node_x, node_y, rr, obstacle_tree):
 
         road_map = []
         n_sample = len(node_x)
-        node_tree = KDTree(np.vstack((node_x, node_y)).T)
+        node_tree = cKDTree(np.vstack((node_x, node_y)).T)
 
         for (i, ix, iy) in zip(range(n_sample), node_x, node_y):
 
-            index, dists = node_tree.search(
-                np.array([ix, iy]).reshape(2, 1), k=n_sample)
+            dists, indexes = node_tree.query([ix, iy], k=n_sample)
 
-            inds = index[0]
             edge_id = []
 
-            for ii in range(1, len(inds)):
-                nx = node_x[inds[ii]]
-                ny = node_y[inds[ii]]
+            for ii in range(1, len(indexes)):
+                nx = node_x[indexes[ii]]
+                ny = node_y[indexes[ii]]
 
                 if not self.is_collision(ix, iy, nx, ny, rr, obstacle_tree):
-                    edge_id.append(inds[ii])
+                    edge_id.append(indexes[ii])
 
                 if len(edge_id) >= self.N_KNN:
                     break
@@ -119,7 +116,7 @@ def voronoi_sampling(sx, sy, gx, gy, ox, oy):
         oxy = np.vstack((ox, oy)).T
 
         # generate voronoi point
-        vor = scipy.spatial.Voronoi(oxy)
+        vor = Voronoi(oxy)
         sample_x = [ix for [ix, _] in vor.vertices]
         sample_y = [iy for [_, iy] in vor.vertices]