diff --git a/tools/annotate_roadmap.py b/tools/annotate_roadmap.py index 2b237e58..c775264c 100644 --- a/tools/annotate_roadmap.py +++ b/tools/annotate_roadmap.py @@ -1,9 +1,13 @@ #!/usr/bin/env python3 import argparse +from multiprocessing import Pool +from typing import List, Tuple import numpy as np import yaml +N_PROCESSES = 8 + if __name__ != "__main__": from . import collision @@ -54,6 +58,12 @@ def add_self_edges(roadmap): return roadmap +def check_proxy(args): + """Prox method to call the collision checker with the right arguments.""" + _, _, E, p0, p1, q0, q1 = args + return collision.ellipsoid_collision_motion(E, p0, p1, q0, q1) + + def compute_edge_conflicts(radius, map): # compute the pairwise collisions and add them to the map E = np.diag([radius, radius]) @@ -61,6 +71,10 @@ def compute_edge_conflicts(radius, map): v_dict = map["roadmap"]["vertices"] edges = map["roadmap"]["edges"] conflicts = [[] for _ in range(num_edges)] + edges_to_check: List[Tuple[ + int, int, np.ndarray, # i, j, E + np.ndarray, np.ndarray, np.ndarray, np.ndarray # p0, p1, q0, q1 + ]] = [] for i in range(0, num_edges): p0 = np.asarray(v_dict[edges[i][0]]) p1 = np.asarray(v_dict[edges[i][1]]) @@ -68,13 +82,17 @@ def compute_edge_conflicts(radius, map): q0 = np.asarray(v_dict[edges[j][0]]) q1 = np.asarray(v_dict[edges[j][1]]) if collision.precheck_bounding_box(E, p0, p1, q0, q1): - collides = collision.ellipsoid_collision_motion( - E, p0, p1, q0, q1) - else: - collides = False - if collides: - conflicts[i].append(j) - conflicts[j].append(i) + edges_to_check.append((i, j, E, p0, p1, q0, q1)) + + # check all edges in parallel + + with Pool(N_PROCESSES) as p: + results = p.map(check_proxy, edges_to_check) + for result, (i, j, _, _, _, _, _) in zip(results, edges_to_check): + if result: + conflicts[i].append(j) + conflicts[j].append(i) + return conflicts