Skip to content

Commit

Permalink
Refactor Basemap=-> BaseMap. Use index_count in render. This will acc…
Browse files Browse the repository at this point in the history 
… the program.
  • Loading branch information
@xuhao1
xuhao1 committed Jul 14, 2022
1 parent 0f148f0 commit ade7af0
Show file tree
Hide file tree
Showing 10 changed files with 88 additions and 92 deletions.
28 changes: 14 additions & 14 deletions scripts/taichislam_node.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,6 +27,8 @@
count = 0

class TaichiSLAMNode:
mesher: MarchingCubeMesher
mapping: BaseMap
def __init__(self):
cuda = rospy.get_param('~use_cuda', True)
self.texture_compressed = rospy.get_param('~texture_compressed', False)
Expand Down Expand Up @@ -180,12 +182,12 @@ def test_mesher(self):
self.mapping.init_sphere()
self.mesher.generate_mesh(1)
mesher = self.mesher
self.render.set_particles(mesher.mesh_vertices, mesher.mesh_vertices)
self.render.set_particles(mesher.mesh_vertices, mesher.mesh_vertices, mesher.num_vetices[None])
self.render.set_mesh(mesher.mesh_vertices, mesher.mesh_colors, mesher.mesh_normals, mesher.mesh_indices)

def update_test_mesher(self):
self.mapping.cvt_TSDF_to_voxels_slice(self.render.slice_z, 100)
self.render.set_particles(self.mapping.export_TSDF_xyz, self.mapping.export_color)
self.render.set_particles(self.mapping.export_TSDF_xyz, self.mapping.export_color, self.mapping.num_TSDF_particles[None])

def process_depth_frame(self, depth_msg, frame):
self.taichimapping_depth_callback(frame, depth_msg)
Expand Down Expand Up @@ -225,10 +227,13 @@ def process_pcl_pose(self, msg, pose):
# self.taichimapping_pcl_callback(pose, xyz_array, rgb_array)

def rendering(self):
mapping = self.mapping
if self.mapping_type == "tsdf" and self.enable_rendering:
if self.render.enable_slice_z:
self.mapping.cvt_TSDF_to_voxels_slice(self.render.slice_z)
self.render.set_particles(self.mapping.export_TSDF_xyz, self.mapping.export_color)
mapping.cvt_TSDF_to_voxels_slice(self.render.slice_z)
else:
mapping.cvt_TSDF_surface_to_voxels()
self.render.set_particles(mapping.export_TSDF_xyz, mapping.export_color, mapping.num_TSDF_particles[None])
self.render.rendering()

def taichimapping_depth_callback(self, frame, depth_msg, texture=np.array([], dtype=int)):
Expand Down Expand Up @@ -287,23 +292,18 @@ def process_taichi(self):
mesher.generate_mesh(1)
t_mesh = (time.time() - start_time)*1000
if self.enable_rendering:
self.render.set_mesh(mesher.mesh_vertices, mesher.mesh_colors, mesher.mesh_normals)
self.render.set_mesh(mesher.mesh_vertices, mesher.mesh_colors, mesher.mesh_normals, mesher.num_vetices[None])
else:
start_time = time.time()
mapping.cvt_TSDF_surface_to_voxels()
t_export = (time.time() - start_time)*1000

par_count = mapping.num_export_TSDF_particles[None]

if self.output_map:
start_time = time.time()
mapping.cvt_TSDF_surface_to_voxels()
t_export = (time.time() - start_time)*1000
par_count = mapping.num_TSDF_particles[None]
start_time = time.time()
self.pub_to_ros(mapping.export_TSDF_xyz.to_numpy()[:par_count],
mapping.export_color.to_numpy()[:par_count], mapping.enable_texture)
t_pubros = (time.time() - start_time)*1000

if self.enable_rendering:
self.render.set_particles(mapping.export_TSDF_xyz, mapping.export_color)

if self.enable_rendering and self.render.lock_pos_drone:
R, T = pose_msg_to_numpy(frame.odom.pose.pose)
self.render.camera_lookat = T
Expand Down
1 change: 1 addition & 0 deletions taichi_slam/mapping/__init__.py
View file
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
from .taichi_octomap import *
from .dense_tsdf import *
from .submap_mapping import *
from .mapping_common import *
from .marching_cube_mesher import *
from .topo_graph import *
26 changes: 13 additions & 13 deletions taichi_slam/mapping/dense_esdf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@

var = [1, 2, 3, 4, 5]
@ti.data_oriented
class DenseSDF(Basemap):
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,
enable_esdf=False, internal_voxels = 10, max_submap_size=1000, is_global_map=False,
Expand Down Expand Up @@ -103,7 +103,7 @@ def initialize_sdf_fields(self):

def initialize_fields(self):
self.num_export_particles = ti.field(dtype=ti.i32, shape=())
self.num_export_TSDF_particles = ti.field(dtype=ti.i32, shape=())
self.num_TSDF_particles = ti.field(dtype=ti.i32, shape=())
self.num_export_ESDF_particles = ti.field(dtype=ti.i32, shape=())

self.export_x = ti.Vector.field(3, dtype=ti.f32, shape=self.max_disp_particles)
Expand Down Expand Up @@ -452,26 +452,26 @@ def cvt_occupy_to_voxels(self):
self.cvt_TSDF_surface_to_voxels()

def cvt_TSDF_surface_to_voxels(self):
self.cvt_TSDF_surface_to_voxels_kernel(self.num_export_TSDF_particles,
self.cvt_TSDF_surface_to_voxels_kernel(self.num_TSDF_particles,
self.export_TSDF_xyz, self.export_color, self.max_disp_particles,
self.clear_last_TSDF_exporting, False)
self.clear_last_TSDF_exporting = False

def cvt_TSDF_surface_to_voxels_to(self, num_export_TSDF_particles, max_disp_particles, export_TSDF_xyz, export_color):
self.cvt_TSDF_surface_to_voxels_kernel(num_export_TSDF_particles,
def cvt_TSDF_surface_to_voxels_to(self, num_TSDF_particles, max_disp_particles, export_TSDF_xyz, export_color):
self.cvt_TSDF_surface_to_voxels_kernel(num_TSDF_particles,
export_TSDF_xyz, export_color, max_disp_particles, False, True)

@ti.kernel
def cvt_TSDF_surface_to_voxels_kernel(self, num_export_TSDF_particles:ti.template(), export_TSDF_xyz:ti.template(),
def cvt_TSDF_surface_to_voxels_kernel(self, num_TSDF_particles:ti.template(), export_TSDF_xyz:ti.template(),
export_color:ti.template(), max_disp_particles:ti.template(), clear_last:ti.template(), add_to_cur:ti.template()):
# Number for TSDF
if clear_last:
for i in range(num_export_TSDF_particles[None]):
for i in range(num_TSDF_particles[None]):
export_color[i] = ti.Vector([0.5, 0.5, 0.5], ti.f32)
export_TSDF_xyz[i] = ti.Vector([-100000, -100000, -100000], ti.f32)

if not add_to_cur:
num_export_TSDF_particles[None] = 0
num_TSDF_particles[None] = 0

disp_floor, disp_ceiling = ti.static(self.disp_floor, self.disp_ceiling)

Expand All @@ -486,8 +486,8 @@ def cvt_TSDF_surface_to_voxels_kernel(self, num_export_TSDF_particles:ti.templat
xyz = self.submap_i_j_k_to_xyz(s, i, j, k)
if xyz[2] > disp_ceiling or xyz[2] < disp_floor:
continue
index = ti.atomic_add(num_export_TSDF_particles[None], 1)
if num_export_TSDF_particles[None] < max_disp_particles:
index = ti.atomic_add(num_TSDF_particles[None], 1)
if num_TSDF_particles[None] < max_disp_particles:
if ti.static(self.enable_texture):
export_color[index] = self.color[s, i, j, k]
export_TSDF_xyz[index] = xyz
Expand All @@ -512,13 +512,13 @@ def cvt_ESDF_to_voxels_slice(self, z: ti.template()):
def cvt_TSDF_to_voxels_slice_kernel(self, z: ti.template(), dz:ti.template()):
_index = int((z+self.map_size_[2]/2.0)/self.voxel_size)
# Number for ESDF
self.num_export_TSDF_particles[None] = 0
self.num_TSDF_particles[None] = 0
for s, i, j, k in self.TSDF:
if s == self.active_submap_id[None]:
if self.TSDF_observed[s, i, j, k] > 0:
if _index - dz < k < _index + dz:
index = ti.atomic_add(self.num_export_TSDF_particles[None], 1)
if self.num_export_TSDF_particles[None] < self.max_disp_particles:
index = ti.atomic_add(self.num_TSDF_particles[None], 1)
if self.num_TSDF_particles[None] < self.max_disp_particles:
self.export_TSDF[index] = self.TSDF[s, i, j, k]
if ti.static(self.is_global_map):
self.export_TSDF_xyz[index] = self.i_j_k_to_xyz(i, j, k)
Expand Down
89 changes: 41 additions & 48 deletions taichi_slam/mapping/dense_tsdf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@

var = [1, 2, 3, 4, 5]
@ti.data_oriented
class DenseTSDF(Basemap):
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,
internal_voxels = 10, max_submap_size=1000, is_global_map=False,
Expand Down Expand Up @@ -40,7 +40,6 @@ def __init__(self, map_scale=[10, 10], voxel_size=0.05, texture_enabled=False, \
self.internal_voxels = internal_voxels
self.max_submap_size = max_submap_size

self.require_clear_exported = False
self.is_global_map = is_global_map
self.disp_ceiling = disp_ceiling
self.disp_floor = disp_floor
Expand Down Expand Up @@ -93,7 +92,7 @@ def initialize_sdf_fields(self):

def initialize_fields(self):
self.num_export_particles = ti.field(dtype=ti.i32, shape=())
self.num_export_TSDF_particles = ti.field(dtype=ti.i32, shape=())
self.num_TSDF_particles = ti.field(dtype=ti.i32, shape=())
self.num_export_ESDF_particles = ti.field(dtype=ti.i32, shape=())

self.export_x = ti.Vector.field(3, dtype=ti.f32, shape=self.max_disp_particles)
Expand Down Expand Up @@ -245,35 +244,52 @@ def process_new_pcl(self):
self.color[xi] = self.new_pcl_sum_color[i, j, k]/c/255.0
self.new_pcl_count[i, j, k] = 0

@ti.kernel
def fuse_submaps_kernel(self, TSDF:ti.template(), W_TSDF:ti.template(), TSDF_observed:ti.template(), color:ti.template()):
for s, i, j, k in TSDF:
if TSDF_observed[s, i, j, k] > 0:
tsdf = TSDF[s, i, j, k]
w_tsdf = W_TSDF[s, i, j, k]
xyz = self.submap_i_j_k_to_xyz(s, i, j, k)
ijk = self.xyz_to_0ijk(xyz)
#Naive merging with weight. TODO: use interpolation
w_new = w_tsdf + self.W_TSDF[ijk]
self.TSDF[ijk] = (self.W_TSDF[ijk]*self.TSDF[ijk] + w_tsdf*tsdf)/w_new
if ti.static(self.enable_texture):
c = color[s, i, j, k]
self.color[ijk] = (self.W_TSDF[ijk]*self.color[ijk] + w_tsdf*c)/w_new
self.W_TSDF[ijk] = w_new
self.TSDF_observed[ijk] = 1

def fuse_submaps(self, submaps):
self.B.parent().deactivate_all()
print("try to fuse all submaps, currently active submap", submaps.active_submap_id[None])
self.fuse_submaps_kernel(submaps.TSDF, submaps.W_TSDF, submaps.TSDF_observed, submaps.color)

def cvt_occupy_to_voxels(self):
self.cvt_TSDF_surface_to_voxels()

def cvt_TSDF_surface_to_voxels(self, add_to_cur = False):
self.cvt_TSDF_surface_to_voxels_kernel(self.num_export_TSDF_particles,
self.export_TSDF_xyz, self.export_color, self.max_disp_particles,
self.require_clear_exported , add_to_cur)
self.require_clear_exported = False
def cvt_TSDF_surface_to_voxels(self):
self.cvt_TSDF_surface_to_voxels_kernel(self.num_TSDF_particles,
self.export_TSDF_xyz, self.export_color, self.max_disp_particles, False)

def cvt_TSDF_surface_to_voxels_to(self, num_export_TSDF_particles, max_disp_particles, export_TSDF_xyz, export_color):
self.cvt_TSDF_surface_to_voxels_kernel(num_export_TSDF_particles,
export_TSDF_xyz, export_color, max_disp_particles, False, True)
def cvt_TSDF_surface_to_voxels_to(self, num_TSDF_particles, max_disp_particles, export_TSDF_xyz, export_color):
self.cvt_TSDF_surface_to_voxels_kernel(num_TSDF_particles,
export_TSDF_xyz, export_color, max_disp_particles, True)

@staticmethod
@ti.func
def clear_last_output(num, export_TSDF_xyz, export_color):
for i in range(num[None]):
export_color[i] = ti.Vector([0.5, 0.5, 0.5], ti.f32)
export_TSDF_xyz[i] = ti.Vector([-100000, -100000, -100000], ti.f32)
num[None] = 0

@ti.kernel
def cvt_TSDF_surface_to_voxels_kernel(self, num_export_TSDF_particles:ti.template(), export_TSDF_xyz:ti.template(),
export_color:ti.template(), max_disp_particles:ti.template(), clear_last:ti.template(), add_to_cur:ti.template()):
# Number for TSDF
if clear_last:
self.clear_last_output(num_export_TSDF_particles, export_TSDF_xyz, export_color)

def cvt_TSDF_surface_to_voxels_kernel(self, num_TSDF_particles:ti.template(), export_TSDF_xyz:ti.template(),
export_color:ti.template(), max_disp_particles:ti.template(), add_to_cur:ti.template()):
if not add_to_cur:
num_export_TSDF_particles[None] = 0
num_TSDF_particles[None] = 0

disp_floor, disp_ceiling = ti.static(self.disp_floor, self.disp_ceiling)

Expand All @@ -288,8 +304,8 @@ def cvt_TSDF_surface_to_voxels_kernel(self, num_export_TSDF_particles:ti.templat
xyz = self.submap_i_j_k_to_xyz(s, i, j, k)
if xyz[2] > disp_ceiling or xyz[2] < disp_floor:
continue
index = ti.atomic_add(num_export_TSDF_particles[None], 1)
if num_export_TSDF_particles[None] < max_disp_particles:
index = ti.atomic_add(num_TSDF_particles[None], 1)
if num_TSDF_particles[None] < max_disp_particles:
if ti.static(self.enable_texture):
export_color[index] = self.color[s, i, j, k]
export_TSDF_xyz[index] = xyz
Expand All @@ -299,49 +315,26 @@ def cvt_TSDF_surface_to_voxels_kernel(self, num_export_TSDF_particles:ti.templat

@ti.kernel
def cvt_TSDF_to_voxels_slice_kernel(self, dz:ti.template(), clear_last:ti.template()):
if clear_last:
self.clear_last_output(self.num_export_TSDF_particles, self.export_TSDF_xyz, self.export_color)
z = self.slice_z[None]
dz = ti.static(dz)
_index = int(z/self.voxel_size)
# Number for ESDF
self.num_export_TSDF_particles[None] = 0
if clear_last:
self.num_TSDF_particles[None] = 0
for s, i, j, k in self.TSDF:
if s == self.active_submap_id[None]:
if self.TSDF_observed[s, i, j, k] > 0:
if _index - dz < k < _index + dz:
index = ti.atomic_add(self.num_export_TSDF_particles[None], 1)
if self.num_export_TSDF_particles[None] < self.max_disp_particles:
index = ti.atomic_add(self.num_TSDF_particles[None], 1)
if self.num_TSDF_particles[None] < self.max_disp_particles:
self.export_TSDF[index] = self.TSDF[s, i, j, k]
if ti.static(self.is_global_map):
self.export_TSDF_xyz[index] = self.i_j_k_to_xyz(i, j, k)
else:
self.export_TSDF_xyz[index] = self.submap_i_j_k_to_xyz(s, i, j, k)
self.export_color[index] = self.color_from_colomap(self.TSDF[s, i, j, k], -0.5, 0.5)

@ti.kernel
def fuse_submaps_kernel(self, TSDF:ti.template(), W_TSDF:ti.template(), TSDF_observed:ti.template(), color:ti.template()):
for s, i, j, k in TSDF:
if TSDF_observed[s, i, j, k] > 0:
tsdf = TSDF[s, i, j, k]
w_tsdf = W_TSDF[s, i, j, k]
xyz = self.submap_i_j_k_to_xyz(s, i, j, k)
ijk = self.xyz_to_0ijk(xyz)
#Naive merging with weight. TODO: use interpolation
w_new = w_tsdf + self.W_TSDF[ijk]
self.TSDF[ijk] = (self.W_TSDF[ijk]*self.TSDF[ijk] + w_tsdf*tsdf)/w_new
if ti.static(self.enable_texture):
c = color[s, i, j, k]
self.color[ijk] = (self.W_TSDF[ijk]*self.color[ijk] + w_tsdf*c)/w_new
self.W_TSDF[ijk] = w_new
self.TSDF_observed[ijk] = 1

def fuse_submaps(self, submaps):
self.B.parent().deactivate_all()
print("try to fuse all submaps, currently active submap", submaps.active_submap_id[None])
self.fuse_submaps_kernel(submaps.TSDF, submaps.W_TSDF, submaps.TSDF_observed, submaps.color)

def cvt_TSDF_to_voxels_slice(self, z, dz=0.5, clear_last=False):
def cvt_TSDF_to_voxels_slice(self, z, dz=0.5, clear_last=True):
self.slice_z[None] = z
self.cvt_TSDF_to_voxels_slice_kernel(dz, clear_last)

Expand Down
2 changes: 1 addition & 1 deletion taichi_slam/mapping/mapping_common.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@ def sign(val):
return (0 < val) - (val < 0)

@ti.data_oriented
class Basemap:
class BaseMap:
def __init__(self, voxel_size_xy, voxel_size_z):
self.input_R = ti.Matrix.field(3, 3, dtype=ti.f32, shape=())
self.input_T = ti.Vector.field(3, dtype=ti.f32, shape=())
Expand Down
19 changes: 9 additions & 10 deletions taichi_slam/mapping/submap_mapping.py
View file
Original file line number Diff line number Diff line change
@@ -1,11 +1,11 @@
from taichi_slam.mapping.mapping_common import Basemap
from taichi_slam.mapping.mapping_common import BaseMap
from .dense_tsdf import DenseTSDF
import time
import numpy as np

class SubmapMapping:
submap_collection: Basemap
global_map: Basemap
submap_collection: BaseMap
global_map: BaseMap
def __init__(self, submap_type=DenseTSDF, keyframe_step=20, sub_opts={}, global_opts={}):
sdf_default_opts = {
'map_scale': [10, 10],
Expand Down Expand Up @@ -66,7 +66,7 @@ def set_export_submap(self, new_submap):
self.export_color = new_submap.export_color
if self.submap_type == DenseTSDF:
self.export_TSDF_xyz = new_submap.export_TSDF_xyz
self.num_export_TSDF_particles = new_submap.num_export_TSDF_particles
self.num_TSDF_particles = new_submap.num_TSDF_particles
else:
self.export_x = new_submap.export_x

Expand Down Expand Up @@ -130,17 +130,16 @@ def recast_depth_to_map(self, R, T, depthmap, texture):
self.frame_count += 1

def cvt_TSDF_to_voxels_slice(self, z):
if len(self.submaps) > 0:
if self.exporting_global:
self.global_map.cvt_TSDF_to_voxels_slice(z)
else:
self.submap_collection.cvt_TSDF_to_voxels_slice(z)
if self.exporting_global:
self.global_map.cvt_TSDF_to_voxels_slice(z)
else:
self.submap_collection.cvt_TSDF_to_voxels_slice(z)

def cvt_TSDF_surface_to_voxels(self):
if len(self.submaps) > 0:
if self.exporting_global:
self.global_map.cvt_TSDF_surface_to_voxels()
self.submap_collection.cvt_TSDF_surface_to_voxels_to(self.global_map.num_export_TSDF_particles,
self.submap_collection.cvt_TSDF_surface_to_voxels_to(self.global_map.num_TSDF_particles,
self.global_map.max_disp_particles, self.export_TSDF_xyz, self.export_color)
else:
self.submap_collection.cvt_TSDF_surface_to_voxels()
Expand Down
2 changes: 1 addition & 1 deletion taichi_slam/mapping/taichi_octomap.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@
import time

@ti.data_oriented
class Octomap(Basemap):
class Octomap(BaseMap):
#If K>2 will be K**3 tree
def __init__(self, map_scale=[10, 10], voxel_size=0.05, min_occupy_thres=3, texture_enabled=False,
min_ray_length=0.3, max_ray_length=3.0, max_disp_particles=1000000, K=2):
Expand Down
2 changes: 1 addition & 1 deletion taichi_slam/mapping/topo_graph.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
from .mapping_common import *
@ti.data_oriented
class TopoGraphGen:
def __init__(self, mapping: Basemap):
def __init__(self, mapping: BaseMap):
self.mapping = mapping
self.generate_uniform_sample_points()

Expand Down
Loading

0 comments on commit ade7af0

Please sign in to comment.