Parallel build skeleton

xuhao1 · Nov 19, 2022 · 14e9f27 · 14e9f27
1 parent 570d437
commit 14e9f27
Show file tree

Hide file tree

Showing 4 changed files with 115 additions and 50 deletions.
diff --git a/scripts/taichislam_node.py b/scripts/taichislam_node.py
@@ -29,6 +29,7 @@ class TaichiSLAMNode:
     mapping: BaseMap
     def __init__(self):
         self.init_params()
+        self.init_topology_generator() #Must multithread beforce init taichi
         if self.cuda:
             ti.init(arch=ti.cuda, dynamic_index=True, offline_cache=True, packed=True, debug=False, device_memory_GB=1.5)
         else:
@@ -86,7 +87,6 @@ def init_params(self):
         self.skeleton_graph_gen_opts['max_raycast_dist'] = rospy.get_param('~skeleton_graph_gen/max_raycast_dist', 2.5)
         self.skeleton_graph_gen_opts['coll_det_num'] = rospy.get_param('~skeleton_graph_gen/coll_det_num', 64)
         self.skeleton_graph_gen_opts['frontier_combine_angle_threshold'] = rospy.get_param('~skeleton_graph_gen/frontier_combine_angle_threshold', 20)
-        print(self.skeleton_graph_gen_opts)
 
     def send_submap_handle(self, buf):
         self.comm.publishBuffer(buf, CHANNEL_SUBMAP)
@@ -219,12 +219,32 @@ def initial_mapping(self):
         self.mapping.set_color_camera_intrinsic(self.Kcolor)
         self.mapping.set_dep_camera_intrinsic(self.Kdep)
 
-        if self.skeleton_graph_gen:
-            print("Initializing skeleton graph generator...")
-            if self.enable_submap:
-                self.topo = TopoGraphGen(self.mapping.global_map, **self.skeleton_graph_gen_opts)
-            else:
-                self.topo = TopoGraphGen(self.mapping, **self.skeleton_graph_gen_opts)
+
+    def init_topology_generator(self):
+        if not self.skeleton_graph_gen:
+            self.topo = None
+        print("Initializing skeleton graph generator thread...")
+        from multiprocessing import Process, Manager
+        from topo_gen_thread import TopoGenThread
+        self.share_map_man = Manager()
+        self.shared_map_d = self.share_map_man.dict()
+        self.shared_map_d["exit"] = False
+        self.shared_map_d["update"] = False
+        self.shared_map_d["topo_graph_viz"] = None
+        params = {
+            "sdf_params": self.get_sdf_opts(),
+            "octo_params": self.get_octo_opts(),
+            "skeleton_graph_gen_opts": self.skeleton_graph_gen_opts,
+            "use_cuda": False #self.cuda
+        }
+        self.topo = Process(target=TopoGenThread, args=[params, self.shared_map_d])
+        self.topo.start()
+
+    def end_topo_thread(self):
+        if self.topo:
+            self.topo.terminate()
+            self.topo.join()
+            self.topo = None
 
     def process_depth_frame(self, depth_msg, frame):
         self.taichimapping_depth_callback(frame, depth_msg)
@@ -397,31 +417,33 @@ def pub_to_ros(self, pos_, colors_, enable_texture):
         else:
             self.pub_occ.publish(point_cloud(pos_, 'world', has_rgb=enable_texture))
 
-    def post_submapfusion_callback(self, global_map):
+    def post_submapfusion_callback(self, global_map: BaseMap):
         self.post_submap_fusion_count += 1
-        if self.post_submap_fusion_count % 10 == 0:
-            start_pt = np.array([1., 0., 0.5])
-            print("start_pt", start_pt)
-            self.topo.reset()
-            s = time.time()
-            num_poly = self.topo.generate_topo_graph(start_pt, max_nodes=10000)
-            print(f"[Topo] Number of polygons: {num_poly} start pt {start_pt} t: {(time.time()-s)*1000:.1f}ms")
-            self.render.set_lines(self.topo.lines_show, self.topo.lines_color, num=self.topo.lines_num[None])
-            self.render.set_mesh(self.topo.tri_vertices, self.topo.tri_colors, mesh_num=self.topo.num_facelets[None])
-
+        if self.topo:
+            #Share the global map with shared memory
+            obj = global_map.export_submap()
+            self.shared_map_d["map_data"]= obj
+            self.shared_map_d['update'] = True
+            # print("Invoking topo skeleton generation")
+            if self.shared_map_d["topo_graph_viz"] is not None:
+                lines = self.shared_map_d["topo_graph_viz"]["lines"]
+                colors = self.shared_map_d["topo_graph_viz"]["colors"]
+                self.render.set_lines(lines, colors)
 
-if __name__ == '__main__':
-
+def slam_main():
     rospy.init_node( 'taichislam_node' )
-
     taichislamnode = TaichiSLAMNode()
     print("TaichiSLAMNode initialized")
-
     rate = rospy.Rate(100) # 100hz
     while not rospy.is_shutdown():
         taichislamnode.process_taichi()
         taichislamnode.handle_comm()
         if taichislamnode.enable_rendering:
             taichislamnode.rendering()
         rate.sleep()
+    taichislamnode.end_topo_thread()
+
+if __name__ == '__main__':
+    slam_main()
+
 
diff --git a/scripts/topo_gen_thread.py b/scripts/topo_gen_thread.py
@@ -0,0 +1,54 @@
+#!/usr/bin/env python3
+from taichi_slam.mapping import DenseTSDF, TopoGraphGen
+import time
+import taichi as ti
+import numpy as np
+class TopoGen:
+    def __init__(self, params_map, params_topo, man_d):
+        self.mapping = DenseTSDF(is_global_map=True, **params_map)
+        self.topo = TopoGraphGen(self.mapping, **params_topo)
+        self.man_d = man_d
+
+    def run(self):
+        print("Start topo graph generation thread")
+        while not self.man_d["exit"]:
+            try:
+                if self.man_d["update"]:
+                    #print all the keys and values
+                    self.loadMap(self.man_d["map_data"])
+                    self.gen_skeleton_graph()
+                    self.man_d["update"] = False
+                time.sleep(1)
+            except KeyboardInterrupt:
+                break
+
+    def loadMap(self, map_data):
+        TSDF = map_data['TSDF']
+        W_TSDF = map_data['W_TSDF']
+        color = map_data['color']
+        indices = map_data['indices']
+        occupy = map_data['occupy']
+        self.mapping.reset()
+        self.mapping.load_numpy(0, indices, TSDF, W_TSDF, occupy, color)
+
+    def gen_skeleton_graph(self):
+        start_pt = np.array([1., 0., 0.5])
+        self.topo.reset()
+        s = time.time()
+        num_nodes = self.topo.generate_topo_graph(start_pt, max_nodes=100000)
+        print(f"[Topo] Number of polygons: {num_nodes} start pt {start_pt} t: {(time.time()-s)*1000:.1f}ms")
+        self.export_topo_graph()
+
+    def export_topo_graph(self):
+        lines = self.topo.lines_show.to_numpy()[0:self.topo.lines_num[None]]
+        colors = self.topo.lines_color.to_numpy()[0:self.topo.lines_num[None]]
+        self.man_d["topo_graph_viz"] = {"lines": lines, "colors": colors}
+
+def TopoGenThread(params, man_d):
+    if params["use_cuda"]:
+        ti.init(arch=ti.cuda, dynamic_index=True, offline_cache=True, packed=True, debug=False)
+    else:
+        ti.init(arch=ti.cpu, dynamic_index=True, offline_cache=True, packed=True, debug=False)
+    print("TopoGenThread: params = ", params, man_d)
+    topo = TopoGen(params["sdf_params"], params["skeleton_graph_gen_opts"], man_d)
+    topo.run()
diff --git a/taichi_slam/mapping/dense_tsdf.py b/taichi_slam/mapping/dense_tsdf.py
@@ -287,9 +287,12 @@ def fuse_submaps_kernel(self, num_submaps: ti.i32, TSDF:ti.template(), W_TSDF:ti
                 self.W_TSDF[ijk] = w_new
                 self.TSDF_observed[ijk] = 1
                 self.occupy[ijk] = self.occupy[ijk] + occ[ijk]
+
+    def reset(self):
+        self.B.parent().deactivate_all()
 
     def fuse_submaps(self, submaps):
-        self.B.parent().deactivate_all()
+        self.reset()
         print("try to fuse all submaps, currently active submap local: ", submaps.active_submap_id[None], 
             " remote: ", submaps.remote_submap_num[None])
         self.fuse_submaps_kernel(submaps.active_submap_id[None], submaps.TSDF, submaps.W_TSDF, 
@@ -456,30 +459,7 @@ def export_submap(self):
         return obj
 
     def saveMap(self, filename):
-        s = time.time()
-        num = self.count_active()
-        indices = np.zeros((num, 3), np.int32)
-        TSDF = np.zeros((num), np.float32)
-        W_TSDF = np.zeros((num), np.float32)
-        occupy = np.zeros((num), np.int32)
-        if self.enable_texture:
-            color = np.zeros((num, 3), np.float32)
-        else:
-            color = np.array([])
-        self.to_numpy(indices, TSDF, W_TSDF, occupy, color)
-        obj = {
-            'indices': indices,
-            'TSDF': TSDF,
-            'W_TSDF': W_TSDF,
-            'color': color,
-            'occupy': occupy,
-            "map_scale": [self.map_size_xy, self.map_size_z],
-            "voxel_size": self.voxel_size,
-            "texture_enabled": self.enable_texture,
-            "num_voxel_per_blk_axis": self.num_voxel_per_blk_axis,
-        }
-        e = time.time()
-        print(f"[SubmapMapping] Saving map to {filename} {num} voxels takes {e-s:.1f} seconds")
+        obj = self.export_submap()
         np.save(filename, obj)
 
     @staticmethod

diff --git a/taichi_slam/utils/visualization.py b/taichi_slam/utils/visualization.py
@@ -45,6 +45,7 @@ def __init__(self, RES_X, RES_Y):
         self.init_grid()
         self.window.show()
         self.lines = None
+        self.lines_color = None
         self.drone_trajs = {}
         self.available_drone_ids = set()
         self.init_drones()
@@ -138,9 +139,17 @@ def set_particles(self, par, color, num=None):
         self.par_num = num
 
     def set_lines(self, lines, color=None, num=None):
-        self.lines = lines
-        self.lines_color = color
-        self.line_vertex_num = num
+        if type(lines) is np.ndarray:
+            #Initialize lines of taichi vector field
+            self.lines_color = ti.Vector.field(3, dtype=ti.f32, shape=lines.shape[0])
+            self.lines = ti.Vector.field(3, dtype=ti.f32, shape=lines.shape[0])
+            self.lines.from_numpy(lines)
+            self.lines_color.from_numpy(color)
+            self.line_vertex_num = num
+        else:
+            self.lines = lines
+            self.lines_color = color
+            self.line_vertex_num = num
 
     def set_mesh(self, mesh, color, normals=None, indices=None, mesh_num=None):
         self.mesh_vertices = mesh