Could build skeleton online

xuhao1 · Nov 19, 2022 · 570d437 · 570d437
1 parent 7d22570
commit 570d437
Show file tree

Hide file tree

Showing 9 changed files with 95 additions and 43 deletions.
diff --git a/data/ri_tsdf.npy b/data/ri_tsdf.npy
diff --git a/launch/taichislam-d435.launch b/launch/taichislam-d435.launch
@@ -14,6 +14,7 @@
     <arg name="enable_submap" default="true" />
     <arg name="drone_id" default="0" />
     <arg name="vins_config" default="/home/xuhao/bags/swarm_raw_parallel2_2021-10-18_17-09/Configs/SwarmConfig1/fisheye_ptgrey_n3/fisheye_cuda.yaml" />
+    <arg name="skeleton_graph_gen" default="false" />
 
     <node pkg="taichislam" type="taichislam_node.py" name="taichislam_node_$(arg drone_id)" output="screen" >
         <param name="use_cuda" value="$(arg use_cuda)" type="boolean" />
@@ -24,6 +25,7 @@
         <param name="enable_submap" value="$(arg enable_submap)" type="boolean" />
         <param name="output_map" value="$(arg output_map)" type="boolean" />
         <param name="drone_id" value="$(arg drone_id)" type="int" />
+        <param name="enable_skeleton_graph_gen" value="$(arg skeleton_graph_gen)" type="boolean" />
 
         <param name="max_ray_length" value="$(arg max_ray_length)" type="double" />
         <param name="min_ray_length" value="$(arg min_ray_length)" type="double" />
@@ -52,13 +54,15 @@
                 fy: 384.2377014160156
                 cx: 323.4873046875
                 cy: 235.0628204345703
-        </rosparam>
-        <rosparam>
             Kcolor:
                 fx: 384.2377014160156
                 fy: 384.2377014160156
                 cx: 323.4873046875
                 cy: 235.0628204345703
+            skeleton_graph_gen:
+                max_raycast_dist: 2.0
+                coll_det_num: 64
+                frontier_combine_angle_threshold: 20
         </rosparam>
         <!-- <rosparam>
             Kcolor:

diff --git a/scripts/taichislam_node.py b/scripts/taichislam_node.py
@@ -32,7 +32,7 @@ def __init__(self):
         if self.cuda:
             ti.init(arch=ti.cuda, dynamic_index=True, offline_cache=True, packed=True, debug=False, device_memory_GB=1.5)
         else:
-            ti.init(arch=ti.cpu, dynamic_index=True, offline_cache=True, packed=True)
+            ti.init(arch=ti.cpu, dynamic_index=True, offline_cache=True, packed=True, debug=False)
         self.disp_level = 0
         self.count = 0
         self.cur_frame = None
@@ -52,6 +52,7 @@ def __init__(self):
         self.initial_mapping()
         self.init_subscribers()
         self.updated_pcl = False
+        self.post_submap_fusion_count = 0
 
     def init_params(self):
         self.cuda = rospy.get_param('~use_cuda', True)
@@ -79,6 +80,13 @@ def init_params(self):
         self.mapping_type = rospy.get_param('~mapping_type', 'tsdf')
         self.texture_enabled = rospy.get_param('~texture_enabled', True)
         self.max_mesh = rospy.get_param('~disp/max_mesh', 1000000)
+
+        self.skeleton_graph_gen = rospy.get_param('~enable_skeleton_graph_gen', False)
+        self.skeleton_graph_gen_opts = {}
+        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)
@@ -189,10 +197,12 @@ def initial_mapping(self):
                 gopts = self.get_octo_opts()
                 subopts = self.get_submap_opts()
                 self.mapping = SubmapMapping(Octomap, global_opts=gopts, sub_opts=subopts, keyframe_step=self.keyframe_step)
+                self.mapping.post_local_to_global_callback = self.post_submapfusion_callback
             else:
                 gopts = self.get_sdf_opts()
                 subopts = self.get_submap_opts()
                 self.mapping = SubmapMapping(DenseTSDF, global_opts=gopts, sub_opts=subopts, keyframe_step=self.keyframe_step)
+                self.mapping.post_local_to_global_callback = self.post_submapfusion_callback
                 if self.enable_mesher:
                     self.mesher = MarchingCubeMesher(self.mapping.submap_collection, self.max_mesh, tsdf_surface_thres=self.voxel_size*5)
             self.mapping.map_send_handle = self.send_submap_handle
@@ -209,6 +219,13 @@ 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 process_depth_frame(self, depth_msg, frame):
         self.taichimapping_depth_callback(frame, depth_msg)
 
@@ -353,6 +370,7 @@ def process_taichi(self):
             return
         self.updated = False
         pose, t_pcl2npy, t_recast = self.recast()
+        self.drone_pose_odom = pose
         if self.enable_rendering:
             self.render.set_drone_pose(0, pose[0], pose[1])
         t_mesh, t_export, t_pubros = self.output(pose[0], pose[1])
@@ -378,6 +396,18 @@ def pub_to_ros(self, pos_, colors_, enable_texture):
             self.pub_occ.publish(point_cloud(pts, 'world', has_rgb=enable_texture))
         else:
             self.pub_occ.publish(point_cloud(pos_, 'world', has_rgb=enable_texture))
+
+    def post_submapfusion_callback(self, global_map):
+        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 __name__ == '__main__':

diff --git a/taichi_slam/mapping/mapping_common.py b/taichi_slam/mapping/mapping_common.py
@@ -156,7 +156,6 @@ def init_colormap(self):
 
     @ti.func
     def raycast(self, pos, dir, max_dist):
-        # print("check dir", dir, " pos ", pos)
         ray_cast_voxels = max_dist/self.voxel_size_[0]
         x_ = ti.Vector([0., 0., 0.], ti.f32)
         succ = False
@@ -165,7 +164,6 @@ def raycast(self, pos, dir, max_dist):
         for _j in range(ray_cast_voxels):
             _len = _j*self.voxel_size
             x_ = dir*_len + pos
-            submap_id = self.active_submap_id[None]
             if self.is_pos_occupy(x_):
                 succ = True
                 break

diff --git a/taichi_slam/mapping/submap_mapping.py b/taichi_slam/mapping/submap_mapping.py
@@ -38,6 +38,7 @@ def __init__(self, submap_type=DenseTSDF, keyframe_step=20, sub_opts={}, global_
         self.last_frame_id = None
         self.active_submap_frame_id = 0
         self.enable_texture = self.global_map.enable_texture
+        self.post_local_to_global_callback = None
         # self.set_exporting_local() # default is exporting local
 
     def create_globalmap(self, global_opts={}):
@@ -89,12 +90,13 @@ def set_frame_poses(self, frame_poses, from_remote=False):
                 T = frame_poses[frame_id][1]
                 self.global_map.set_base_pose_submap(self.submaps[frame_id], R, T)
                 used_poses[frame_id] = frame_poses[frame_id]
-        print(f"[SubmapMapping] Update frame poses from PGO cost {(time.time() - s)*1000:.1f}ms")
+        # print(f"[SubmapMapping] Update frame poses from PGO cost {(time.time() - s)*1000:.1f}ms")
         #We need to broadcast the poses to remote
         if not from_remote:
             self.send_traj(used_poses)
 
     def create_new_submap(self, frame_id, R, T):
+        print("[SubmapMapping] Create new submap ", frame_id)
         if self.first_init:
             self.first_init = False
         else:
@@ -111,8 +113,8 @@ def create_new_submap(self, frame_id, R, T):
         self.active_submap_frame_id = frame_id
 
         print(f"[SubmapMapping] Created new submap on frame {frame_id}, now have {submap_id+1} submaps")
-        # if submap_id % 2 == 0:
-        #     self.saveMap("/home/xuhao/output/test_map.npy")
+        if submap_id % 2 == 0:
+            self.saveMap("/home/xuhao/output/test_map.npy")
         return self.submap_collection
 
     def need_create_new_submap(self, is_keyframe, R, T):
@@ -124,11 +126,10 @@ def need_create_new_submap(self, is_keyframe, R, T):
             return True
         return False
 
-    def recast_pcl_to_map(self, R, T, xyz_array, rgb_array, n):
-        pass
-
     def local_to_global(self):
         self.global_map.fuse_submaps(self.submap_collection)
+        if self.post_local_to_global_callback is not None:
+            self.post_local_to_global_callback(self.global_map)
 
     def convert_by_pgo(self, frame_id, R, T):
         self.ego_motion_poses[frame_id] = (R, T)
@@ -140,11 +141,11 @@ def convert_by_pgo(self, frame_id, R, T):
         return R, T
 
     def recast_depth_to_map_by_frame(self, frame_id, is_keyframe, pose, ext, depthmap, texture):
+        # print("[SubmapMapping] Recast depth to map by frame ", frame_id)
         R, T = pose
         R_ext, T_ext = ext
         R, T = self.convert_by_pgo(frame_id, R, T)
         if self.need_create_new_submap(is_keyframe, R, T):
-            #In early debug we use framecount as frameid
             self.create_new_submap(frame_id, R, T)
         Rcam = R @ R_ext
         Tcam = T + R @ T_ext

diff --git a/taichi_slam/mapping/topo_graph.py b/taichi_slam/mapping/topo_graph.py
@@ -187,6 +187,17 @@ def init_fields(self, max_facelets, coll_det_num, facelet_nh_search_queue_size=1
         self.lines_num = ti.field(dtype=ti.i32, shape=())
         self.lines_num[None] = 0
 
+    def reset(self):
+        self.num_facelets[None] = 0
+        self.num_polyhedron[None] = 0
+        self.num_frontiers[None] = 0
+        self.search_frontiers_idx[None] = 0
+        self.facelet_nh_search_idx[None] = 0
+        self.facelet_nh_search_queue_size[None] = 0
+        self.lines_num[None] = 0
+        self.black_num[None] = 0
+        self.white_num[None] = 0
+
     def generate_uniform_sample_points(self, npoints):
         phi = np.pi * (3 - np.sqrt(5))
         ret = []
@@ -245,9 +256,9 @@ def generate_topo_graph(self, start_pt, max_nodes=100, show=False):
         while self.search_frontiers_idx[None] < self.num_frontiers[None] and self.search_frontiers_idx[None] < max_nodes:
             if self.verify_frontier(self.search_frontiers_idx[None]):
                 frontier = self.frontiers[self.search_frontiers_idx[None]]
-                # print(f"expansion from frontier {self.search_frontiers_idx[None]}, pos :", frontier.next_node_initial)
                 self.node_expansion(frontier.next_node_initial, show, last_node_idx=frontier.master_idx)
             self.search_frontiers_idx[None] += 1
+        return self.num_polyhedron[None]
 
     def generate_mesh_from_hull(self, hull, start_pt):
         if not isinstance(start_pt, np.ndarray):
@@ -331,8 +342,6 @@ def construct_frontier(self, node_idx, idx_start_facelet):
 
     @ti.kernel
     def add_mesh(self, mesh: ti.types.ndarray(), neighbors:ti.types.ndarray(), last_node_idx: ti.i32):
-        index_poly = self.num_polyhedron[None]
-        idx_start_facelet = self.num_facelets[None]
         num_facelets = mesh.shape[0]
         facelet_start_idx = ti.atomic_add(self.num_facelets[None], num_facelets)
         center_pos = ti.Vector([0., 0., 0.], ti.f32)
@@ -348,25 +357,21 @@ def add_mesh(self, mesh: ti.types.ndarray(), neighbors:ti.types.ndarray(), last_
             center_pos += v0 + v1 + v2
             center_count += 3.0
             naive_norm = (v0 + v1 + v2 - 3.0*self.start_point[None]).normalized()
-            self.facelets[facelet_idx].init(index_poly, facelet_idx, v0, v1, v2, naive_norm)
+            self.facelets[facelet_idx].init(self.num_polyhedron[None], facelet_idx, v0, v1, v2, naive_norm)
             self.facelets[facelet_idx].is_frontier = self.detect_facelet_frontier(self.facelets[facelet_idx])
             #For visualization
             for j in range(ti.static(3)):
-                self.tri_colors[facelet_idx*3 + j] = self.colormap[index_poly]
+                self.tri_colors[facelet_idx*3 + j] = self.colormap[self.num_polyhedron[None]]
                 if self.facelets[facelet_idx].is_frontier:
                     self.tri_colors[facelet_idx*3 + j][3] = ti.static(self.transparent_frontier)
-            # if self.facelets[facelet_idx].is_frontier:
-            #     self.tri_vertices[facelet_idx*3] = ti.Vector([0.0, 0.0, 0.0])
-            #     self.tri_vertices[facelet_idx*3+1] = ti.Vector([0.0, 0.0, 0.0])
-            #     self.tri_vertices[facelet_idx*3+2] = ti.Vector([0.0, 0.0, 0.0])
         new_node_center = center_pos/center_count
-        self.nodes[index_poly].init(index_poly, last_node_idx, facelet_start_idx, facelet_start_idx + num_facelets, center_pos/center_count)
+        self.nodes[self.num_polyhedron[None]].init(self.num_polyhedron[None], last_node_idx, facelet_start_idx, facelet_start_idx + num_facelets, center_pos/center_count)
         if last_node_idx >= 0:
             self.add_line(self.nodes[last_node_idx].center, new_node_center, ti.Vector([0.0, 0.0, 0.0]), ti.Vector([0.0, 0.0, 0.0]))
         #Construct facelet from neighbors relationship
         ti.loop_config(serialize=True)
-        for i in range(idx_start_facelet, idx_start_facelet+num_facelets):
-            idx = i - idx_start_facelet  #This idx is define in hull
+        for i in range(facelet_start_idx, facelet_start_idx+num_facelets):
+            idx = i - facelet_start_idx  #This idx is define in hull
             if not self.facelets[i].assigned and self.facelets[i].is_frontier:
                 # print("\nfacelet ", self.facelets[i].facelet_idx, "idx", idx, " not assigned, use as seed")
                 self.facelet_nh_search_idx[None] = 0
@@ -376,19 +381,19 @@ def add_mesh(self, mesh: ti.types.ndarray(), neighbors:ti.types.ndarray(), last_
                 while self.facelet_nh_search_idx[None] < self.facelet_nh_search_queue_size[None]:
                     search_idx = ti.atomic_add(self.facelet_nh_search_idx[None], 1)
                     _idx = self.facelet_nh_search_queue[search_idx] #This idx is define in hull
-                    facelet_idx = _idx + idx_start_facelet
+                    facelet_idx = _idx + facelet_start_idx
                     self.facelets[facelet_idx].assigned = True
                     for j in range(ti.static(3)):
                         #Check if neighor is assigned and frontier
-                        idx_neighbor = neighbors[_idx, j] + idx_start_facelet
+                        idx_neighbor = neighbors[_idx, j] + facelet_start_idx
                         if self.facelets[idx_neighbor].is_frontier and not self.facelets[idx_neighbor].assigned and \
                                 normal.dot(self.facelets[idx_neighbor].normal) > ti.static(self.frontier_normal_dot_threshold):
                             self.facelet_nh_search_queue[ti.atomic_add(self.facelet_nh_search_queue_size[None], 1)] = neighbors[_idx, j]
                             # print("neighbor", neighbors[_idx, j], "idx", idx_neighbor, "v0", self.facelets[idx_neighbor].v0, "v1", self.facelets[idx_neighbor].v1, "v2", self.facelets[idx_neighbor].v2)
                 # Now we can construct the frontier from these facelets
                 # It uses logs in facelet_nh_search_queue to find the facelets
-                self.construct_frontier(index_poly, idx_start_facelet)
-        self.num_polyhedron[None] = index_poly + 1
+                self.construct_frontier(self.num_polyhedron[None], facelet_start_idx)
+        self.num_polyhedron[None] = self.num_polyhedron[None] + 1
 
     @ti.kernel
     def detect_collisions(self) ->ti.i32:
@@ -400,7 +405,6 @@ def detect_collisions(self) ->ti.i32:
         for i in range(ti.static(self.coll_det_num)):
             succ, t, col_pos, _len = self.raycast(pos, self.sample_dirs[i], max_raycast_dist)
             if succ:
-                # print("is collision on dir", self.sample_dirs[i], " pt ", col_pos)
                 index = ti.atomic_add(self.black_num[None], 1)
                 self.black_list[index] = col_pos
                 self.black_unit_list[index] = self.sample_dirs[i]

diff --git a/taichi_slam/utils/communication.py b/taichi_slam/utils/communication.py
@@ -23,20 +23,20 @@ def publishBuffer(self, buf, channel=CHANNEL_SUBMAP):
         msg.buffer = buf
         self.sent_msgs.add(msg.msg_id)
         self.lcm.publish(channel, msg.encode())
-        print(f"Sent message on channel {channel} msg_id {msg.msg_id} size {len(msg.buffer)/1024:.1f} KB")
+        # print(f"Sent message on channel {channel} msg_id {msg.msg_id} size {len(msg.buffer)/1024:.1f} KB")
 
     def handle_submap(self, channel, data):
         msg = Buffer.decode(data)
         if msg.msg_id in self.sent_msgs:
             return
-        print(f"Received message on channel {channel} msg_id {msg.msg_id}")
+        # print(f"Received message on channel {channel} msg_id {msg.msg_id}")
         self.on_submap(msg.buffer)
 
     def handle_traj(self, channel, data):
         msg = Buffer.decode(data)
         if msg.msg_id in self.sent_msgs:
             return
-        print(f"Received message on channel {channel} msg_id {msg.msg_id}")
+        # print(f"Received message on channel {channel} msg_id {msg.msg_id}")
         self.on_traj(msg.buffer)
 
     def handle(self):

diff --git a/taichi_slam/utils/visualization.py b/taichi_slam/utils/visualization.py
@@ -200,8 +200,8 @@ def rendering(self):
                index_count=self.mesh_num, per_vertex_color=self.mesh_color, two_sided=True)
 
         # #Some additional lines        
-        # if self.lines is not None:
-        #     scene.lines(self.lines, self.grid_width*5, per_vertex_color=self.lines_color, vertex_count=self.line_vertex_num)
+        if self.lines is not None:
+            scene.lines(self.lines, self.grid_width*5, per_vertex_color=self.lines_color, vertex_count=self.line_vertex_num)
 
         #Drone frame
         scene.lines(self.drone_frame_lines, self.drone_frame_line_width, per_vertex_color=self.drone_frame_colors, vertex_count=self.drone_num()*6)