export submap to buffer

xuhao1 · Oct 28, 2022 · 7cd80ef · 7cd80ef
1 parent 265b3ca
commit 7cd80ef
Show file tree

Hide file tree

Showing 6 changed files with 92 additions and 48 deletions.
diff --git a/TaichiSLAM_demo.py b/TaichiSLAM_demo.py
@@ -149,7 +149,7 @@ def ros_subscribe_pcl():
             min_occupy_thres = args.occupy_thres,
             map_scale=args.map_size,
             voxel_size=args.voxel_size,
-            block_size=args.blk)
+            num_voxel_per_blk_axis=args.blk)
 
     scene.init_control(gui, radius=6, theta=-math.pi/4,center=(0, 0, 0), is_ortho=True)
 

diff --git a/scripts/taichislam_node.py b/scripts/taichislam_node.py
@@ -140,7 +140,7 @@ def get_octo_opts(self):
     def get_sdf_opts(self):
         opts = self.get_general_mapping_opts()
         opts.update({
-            'block_size': rospy.get_param('~block_size', 16),  #How many voxels per block per axis
+            'num_voxel_per_blk_axis': rospy.get_param('~num_voxel_per_blk_axis', 16),  #How many voxels per block per axis
         })
         return opts
 

diff --git a/taichi_slam/mapping/dense_esdf.py b/taichi_slam/mapping/dense_esdf.py
@@ -10,21 +10,21 @@
 @ti.data_oriented
 class DenseSDF(BaseMap):
     def __init__(self, map_scale=[10, 10], voxel_size=0.05, min_occupy_thres=0, texture_enabled=False, \
-            max_disp_particles=1000000, block_size=16, max_ray_length=10, min_ray_length=0.3, 
+            max_disp_particles=1000000, num_voxel_per_blk_axis=16, max_ray_length=10, min_ray_length=0.3, 
             enable_esdf=False, internal_voxels = 10, max_submap_size=1000, is_global_map=False, 
             disp_ceiling=1.8, disp_floor=-0.3):
         super(DenseSDF, self).__init__()
         self.map_size_xy = map_scale[0]
         self.map_size_z = map_scale[1]
 
-        self.block_size = block_size
+        self.num_voxel_per_blk_axis = num_voxel_per_blk_axis
         self.voxel_size = voxel_size
 
-        self.N = math.ceil(map_scale[0] / voxel_size/block_size)*block_size
-        self.Nz = math.ceil(map_scale[1] / voxel_size/block_size)*block_size
+        self.N = math.ceil(map_scale[0] / voxel_size/num_voxel_per_blk_axis)*num_voxel_per_blk_axis
+        self.Nz = math.ceil(map_scale[1] / voxel_size/num_voxel_per_blk_axis)*num_voxel_per_blk_axis
 
-        self.block_num_xy = math.ceil(map_scale[0] / voxel_size/block_size)
-        self.block_num_z = math.ceil(map_scale[1] / voxel_size/block_size)
+        self.block_num_xy = math.ceil(map_scale[0] / voxel_size/num_voxel_per_blk_axis)
+        self.block_num_z = math.ceil(map_scale[1] / voxel_size/num_voxel_per_blk_axis)
 
         self.map_size_xy = voxel_size * self.N
         self.map_size_z = voxel_size * self.Nz
@@ -49,22 +49,22 @@ def __init__(self, map_scale=[10, 10], voxel_size=0.05, min_occupy_thres=0, text
 
         self.initialize_fields()
 
-    def data_structures(self, submap_num, block_num_xy, block_num_z, block_size):
-        if block_size < 1:
-            print("block_size must be greater than 1")
+    def data_structures(self, submap_num, block_num_xy, block_num_z, num_voxel_per_blk_axis):
+        if num_voxel_per_blk_axis < 1:
+            print("num_voxel_per_blk_axis must be greater than 1")
             exit(0)
         if self.is_global_map:
             Broot = ti.root.pointer(ti.ijkl, (1, block_num_xy, block_num_xy, block_num_z))
-            B = Broot.dense(ti.ijkl, (1, block_size, block_size, block_size))
+            B = Broot.dense(ti.ijkl, (1, num_voxel_per_blk_axis, num_voxel_per_blk_axis, num_voxel_per_blk_axis))
         else:
             Broot = ti.root.pointer(ti.i, submap_num)
-            B = Broot.pointer(ti.ijkl, (1, block_num_xy, block_num_xy, block_num_z)).dense(ti.ijkl, (1, block_size, block_size, block_size))
+            B = Broot.pointer(ti.ijkl, (1, block_num_xy, block_num_xy, block_num_z)).dense(ti.ijkl, (1, num_voxel_per_blk_axis, num_voxel_per_blk_axis, num_voxel_per_blk_axis))
         return B, Broot
 
-    def data_structures_grouped(self, block_num_xy, block_num_z, block_size):
-        if block_size > 1:
+    def data_structures_grouped(self, block_num_xy, block_num_z, num_voxel_per_blk_axis):
+        if num_voxel_per_blk_axis > 1:
             Broot = ti.root.pointer(ti.ijk, (block_num_xy, block_num_xy, block_num_z))
-            B = Broot.dense(ti.ijk, (block_size, block_size, block_size))
+            B = Broot.dense(ti.ijk, (num_voxel_per_blk_axis, num_voxel_per_blk_axis, num_voxel_per_blk_axis))
         else:
             B = ti.root.dense(ti.ijk, (block_num_xy, block_num_xy, block_num_z))
             Broot = B
@@ -73,7 +73,7 @@ def data_structures_grouped(self, block_num_xy, block_num_z, block_size):
     def initialize_sdf_fields(self):
         block_num_xy = self.block_num_xy
         block_num_z = self.block_num_z
-        block_size = self.block_size
+        num_voxel_per_blk_axis = self.num_voxel_per_blk_axis
         submap_num = self.max_submap_size
         if self.is_global_map:
             submap_num = 1
@@ -90,15 +90,15 @@ def initialize_sdf_fields(self):
             self.observed = ti.field(dtype=ti.i8)
             self.fixed = ti.field(dtype=ti.i8)
             self.parent_dir = ti.Vector.field(3, dtype=ti.i32)
-        self.B, self.Broot = self.data_structures(submap_num, block_num_xy, block_num_z, block_size)
+        self.B, self.Broot = self.data_structures(submap_num, block_num_xy, block_num_z, num_voxel_per_blk_axis)
         self.B.place(self.W_TSDF,self.TSDF, self.TSDF_observed)
         if self.enable_esdf:
             self.B.place(self.ESDF, self.observed, self.fixed, self.parent_dir)
         if self.enable_texture:
             self.B.place(self.color)
         if self.enable_esdf:
             self.updated_TSDF = ti.field(dtype=ti.i32)
-            self.T, self.Troot = self.data_structures(submap_num, block_num_xy, block_num_z, block_size)
+            self.T, self.Troot = self.data_structures(submap_num, block_num_xy, block_num_z, num_voxel_per_blk_axis)
             self.T.place(self.updated_TSDF)
 
     def initialize_fields(self):
@@ -121,7 +121,7 @@ def initialize_fields(self):
         self.new_pcl_count = ti.field(dtype=ti.i32)
         self.new_pcl_sum_pos = ti.Vector.field(3, dtype=ti.f32) #position in sensor coor
         self.new_pcl_z = ti.field(dtype=ti.f32) #position in sensor coor
-        self.PCL, self.PCLroot = self.data_structures_grouped(self.block_num_xy, self.block_num_z, self.block_size)
+        self.PCL, self.PCLroot = self.data_structures_grouped(self.block_num_xy, self.block_num_z, self.num_voxel_per_blk_axis)
         self.PCL.place(self.new_pcl_count, self.new_pcl_sum_pos, self.new_pcl_z)
 
         self.initialize_sdf_fields()

diff --git a/taichi_slam/mapping/dense_tsdf.py b/taichi_slam/mapping/dense_tsdf.py
@@ -11,21 +11,21 @@
 @ti.data_oriented
 class DenseTSDF(BaseMap):
     def __init__(self, map_scale=[10, 10], voxel_size=0.05, texture_enabled=False, \
-            max_disp_particles=1000000, block_size=16, max_ray_length=10, min_ray_length=0.3,
+            max_disp_particles=1000000, num_voxel_per_blk_axis=16, max_ray_length=10, min_ray_length=0.3,
             internal_voxels = 10, max_submap_size=1000, is_global_map=False, 
             disp_ceiling=1.8, disp_floor=-0.3):
         super(DenseTSDF, self).__init__(voxel_size)
         self.map_size_xy = map_scale[0]
         self.map_size_z = map_scale[1]
 
-        self.block_size = block_size
+        self.num_voxel_per_blk_axis = num_voxel_per_blk_axis
         self.voxel_size = voxel_size
 
-        self.N = math.ceil(map_scale[0] / voxel_size/block_size)*block_size
-        self.Nz = math.ceil(map_scale[1] / voxel_size/block_size)*block_size
+        self.N = math.ceil(map_scale[0] / voxel_size/num_voxel_per_blk_axis)*num_voxel_per_blk_axis
+        self.Nz = math.ceil(map_scale[1] / voxel_size/num_voxel_per_blk_axis)*num_voxel_per_blk_axis
 
-        self.block_num_xy = math.ceil(map_scale[0] / voxel_size/block_size)
-        self.block_num_z = math.ceil(map_scale[1] / voxel_size/block_size)
+        self.block_num_xy = math.ceil(map_scale[0] / voxel_size/num_voxel_per_blk_axis)
+        self.block_num_z = math.ceil(map_scale[1] / voxel_size/num_voxel_per_blk_axis)
 
         self.map_size_xy = voxel_size * self.N
         self.map_size_z = voxel_size * self.Nz
@@ -47,22 +47,22 @@ def __init__(self, map_scale=[10, 10], voxel_size=0.05, texture_enabled=False, \
         self.initialize_fields()
         print(f"TSDF map initialized blocks {self.block_num_xy}x{self.block_num_xy}x{self.block_num_z}")
 
-    def data_structures(self, submap_num, block_num_xy, block_num_z, block_size):
-        if block_size < 1:
-            print("block_size must be greater than 1")
+    def data_structures(self, submap_num, block_num_xy, block_num_z, num_voxel_per_blk_axis):
+        if num_voxel_per_blk_axis < 1:
+            print("num_voxel_per_blk_axis must be greater than 1")
             exit(0)
         if self.is_global_map:
             Broot = ti.root.pointer(ti.ijkl, (1, block_num_xy, block_num_xy, block_num_z))
-            B = Broot.dense(ti.ijkl, (1, block_size, block_size, block_size))
+            B = Broot.dense(ti.ijkl, (1, num_voxel_per_blk_axis, num_voxel_per_blk_axis, num_voxel_per_blk_axis))
         else:
-            Broot = ti.root.pointer(ti.i, submap_num)
-            B = Broot.pointer(ti.ijkl, (1, block_num_xy, block_num_xy, block_num_z)).dense(ti.ijkl, (1, block_size, block_size, block_size))
+            Broot = ti.root.pointer(ti.i, submap_num).pointer(ti.ijkl, (1, block_num_xy, block_num_xy, block_num_z))
+            B = Broot.dense(ti.ijkl, (1, num_voxel_per_blk_axis, num_voxel_per_blk_axis, num_voxel_per_blk_axis))
         return B, Broot
 
-    def data_structures_grouped(self, block_num_xy, block_num_z, block_size):
-        if block_size > 1:
+    def data_structures_grouped(self, block_num_xy, block_num_z, num_voxel_per_blk_axis):
+        if num_voxel_per_blk_axis > 1:
             Broot = ti.root.pointer(ti.ijk, (block_num_xy, block_num_xy, block_num_z))
-            B = Broot.dense(ti.ijk, (block_size, block_size, block_size))
+            B = Broot.dense(ti.ijk, (num_voxel_per_blk_axis, num_voxel_per_blk_axis, num_voxel_per_blk_axis))
         else:
             B = ti.root.dense(ti.ijk, (block_num_xy, block_num_xy, block_num_z))
             Broot = B
@@ -71,7 +71,7 @@ def data_structures_grouped(self, block_num_xy, block_num_z, block_size):
     def initialize_sdf_fields(self):
         block_num_xy = self.block_num_xy
         block_num_z = self.block_num_z
-        block_size = self.block_size
+        num_voxel_per_blk_axis = self.num_voxel_per_blk_axis
         submap_num = self.max_submap_size
         if self.is_global_map:
             submap_num = 1
@@ -86,7 +86,7 @@ def initialize_sdf_fields(self):
             self.color = ti.Vector.field(3, dtype=ti.f16)
         else:
             self.color = None
-        self.B, self.Broot = self.data_structures(submap_num, block_num_xy, block_num_z, block_size)
+        self.B, self.Broot = self.data_structures(submap_num, block_num_xy, block_num_z, num_voxel_per_blk_axis)
         self.B.place(self.W_TSDF,self.TSDF, self.TSDF_observed, self.occupy, offset=offset)
         if self.enable_texture:
             self.B.place(self.color, offset=offset)
@@ -109,9 +109,9 @@ def initialize_fields(self):
         self.new_pcl_count = ti.field(dtype=ti.i32)
         self.new_pcl_sum_pos = ti.Vector.field(3, dtype=ti.f16) #position in sensor coor
         self.new_pcl_z = ti.field(dtype=ti.f16) #position in sensor coor
-        grp_block_num = max(int(3.2*self.max_ray_length/self.block_size/self.voxel_size), 1)
-        self.PCL, self.PCLroot = self.data_structures_grouped(grp_block_num, grp_block_num, self.block_size)
-        offset = [-self.block_size*grp_block_num//2, -self.block_size*grp_block_num//2, -self.block_size*grp_block_num//2]
+        grp_block_num = max(int(3.2*self.max_ray_length/self.num_voxel_per_blk_axis/self.voxel_size), 1)
+        self.PCL, self.PCLroot = self.data_structures_grouped(grp_block_num, grp_block_num, self.num_voxel_per_blk_axis)
+        offset = [-self.num_voxel_per_blk_axis*grp_block_num//2, -self.num_voxel_per_blk_axis*grp_block_num//2, -self.num_voxel_per_blk_axis*grp_block_num//2]
         self.PCL.place(self.new_pcl_count, self.new_pcl_sum_pos, self.new_pcl_z, offset=offset)
 
         self.slice_z = ti.field(dtype=ti.f16, shape=())
@@ -385,7 +385,7 @@ def get_voxels_TSDF_slice(self, z):
     def finalization_current_submap(self):
         count = self.count_active_func()/1024
         count_mem = count * ti.static(self.mem_per_voxel)/1024
-        print(f"Will finalize submap {self.active_submap_id[None]} opened voxel: {count}k,{count_mem}MB")
+        # print(f"Will finalize submap {self.active_submap_id[None]} opened voxel: {count}k,{count_mem}MB")
 
     @ti.func
     def count_active_func(self):
@@ -412,8 +412,8 @@ def to_numpy(self, data_indices: ti.types.ndarray(element_dim=1), data_tsdf: ti.
                     data_tsdf[_count] = self.TSDF[s, i, j, k]
                     data_wtsdf[_count] = self.W_TSDF[s, i, j, k]
                     data_occ[_count] = self.occupy[s, i, j, k]
-                    # if ti.static(self.enable_texture):
-                    #     data_color[_count] = self.color[s, i, j, k]
+                    if ti.static(self.enable_texture):
+                        data_color[_count] = self.color[s, i, j, k]
 
     @ti.kernel
     def load_numpy(self, data_indices: ti.types.ndarray(element_dim=1), data_tsdf: ti.types.ndarray(), data_wtsdf: ti.types.ndarray(), data_occ: ti.types.ndarray(), data_color:ti.types.ndarray()):
@@ -427,6 +427,32 @@ def load_numpy(self, data_indices: ti.types.ndarray(element_dim=1), data_tsdf: t
                 self.color[sijk] = data_color[i]
             self.TSDF_observed[sijk] = 1
 
+    def export_submap(self):
+        s = time.time()
+        num = self.count_active()
+        indices = np.zeros((num, 3), np.int16)
+        tsdf = np.zeros((num), np.float16)
+        w_tsdf = np.zeros((num), np.float16)
+        occupy = np.zeros((num), np.int8)
+        if self.enable_texture:
+            color = np.zeros((num, 3), np.float16)
+        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,
+        }
+        print(f"Export submap {self.active_submap_id[None]} to numpy, voxels {num/1024:.1f}k, time: {1000*(time.time()-s):.1f}ms")
+        return obj
+
     def saveMap(self, filename):
         s = time.time()
         num = self.count_active()
@@ -448,7 +474,7 @@ def saveMap(self, filename):
             "map_scale": [self.map_size_xy, self.map_size_z],
             "voxel_size": self.voxel_size,
             "texture_enabled": self.enable_texture,
-            "block_size": self.block_size,
+            "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")
@@ -463,7 +489,7 @@ def loadMap(filename):
         indices = obj['indices']
         occupy = obj['occupy']
         mapping = DenseTSDF(map_scale=obj['map_scale'], voxel_size=obj['voxel_size'], 
-            texture_enabled=obj['texture_enabled'], block_size=obj['block_size'], is_global_map=True)
+            texture_enabled=obj['texture_enabled'], num_voxel_per_blk_axis=obj['num_voxel_per_blk_axis'], is_global_map=True)
         mapping.load_numpy(indices, TSDF, W_TSDF, occupy, color)
         print(f"[SubmapMapping] Loaded {TSDF.shape[0]} voxels from {filename}")
         return mapping
diff --git a/taichi_slam/mapping/submap_mapping.py b/taichi_slam/mapping/submap_mapping.py
@@ -2,6 +2,8 @@
 from .dense_tsdf import DenseTSDF
 import time
 import numpy as np
+import io
+import zlib
 
 class SubmapMapping:
     submap_collection: BaseMap
@@ -14,7 +16,7 @@ def __init__(self, submap_type=DenseTSDF, keyframe_step=20, sub_opts={}, global_
             'min_ray_length': 0.3,
             'max_ray_length': 3.0,
             'max_disp_particles': 100000,
-            'block_size': 10,
+            'num_voxel_per_blk_axis': 10,
             'max_submap_size': 1000
         }
         sdf_default_opts.update(sub_opts)
@@ -34,6 +36,7 @@ def __init__(self, submap_type=DenseTSDF, keyframe_step=20, sub_opts={}, global_
         self.ego_motion_poses = {}
         self.pgo_poses = {}
         self.last_frame_id = None
+        self.active_submap_frame_id = 0
         # self.set_exporting_local() # default is exporting local
 
     def create_globalmap(self, global_opts={}):
@@ -44,7 +47,7 @@ def create_globalmap(self, global_opts={}):
             'min_ray_length': 0.3,
             'max_ray_length': 3.0,
             'max_disp_particles': 1000000,
-            'block_size': 10,
+            'num_voxel_per_blk_axis': 10,
             'max_submap_size': 1000,
             'is_global_map': True
         }
@@ -84,11 +87,22 @@ def set_frame_poses(self, frame_poses):
                 T = frame_poses[frame_id][1]
                 self.global_map.set_base_pose_submap(self.submaps[frame_id], R, T)
         print(f"[SubmapMapping] Update frame poses from PGO cost {(time.time() - s)*1000:.1f}ms")
+
+    def send_submap(self, submap):
+        submap["frame_id"] = self.active_submap_frame_id
+        f = io.BytesIO()
+        np.save(f, submap)
+        s = time.time()
+        compressed = zlib.compress(f.getbuffer(), level=1)
+        print(f"[SubmapMapping] Send submap with {len(f.getbuffer())/1024:.1f} kB, compressed {len(compressed)/1024:.1f}kB compress cost {(time.time() - s)*1000:.1f}ms")
+
 
     def create_new_submap(self, frame_id, R, T):
         if self.first_init:
             self.first_init = False
         else:
+            submap = self.submap_collection.export_submap()
+            self.send_submap(submap)
             self.submap_collection.switch_to_next_submap()
             self.submap_collection.clear_last_TSDF_exporting = True
             self.local_to_global()
@@ -97,6 +111,7 @@ def create_new_submap(self, frame_id, R, T):
         self.submap_collection.set_base_pose_submap(submap_id, R, T)
         self.submap_collection.set_base_pose_submap_kernel(submap_id, R, T)
         self.submaps[frame_id] = submap_id
+        self.active_submap_frame_id = frame_id
 
         print(f"[SubmapMapping] Created new submap, now have {submap_id+1} submaps")
         # if submap_id % 2 == 0:
@@ -175,3 +190,6 @@ def cvt_TSDF_surface_to_voxels(self):
 
     def saveMap(self, filename):
         self.global_map.saveMap(filename)
+
+    def export_submap(self):
+        self.submap_collection.export_submap()
diff --git a/tests/marching_cube_test.py b/tests/marching_cube_test.py
@@ -12,7 +12,7 @@
                 min_occupy_thres = 1,
                 map_scale=[100, 100],
                 voxel_size=0.05,
-                block_size=16,
+                num_voxel_per_blk_axis=16,
                 enable_esdf=False,
                 max_ray_length=10)
     max_mesh = 1000000