Skip to content

Commit 5b418a9

Browse files
xuhao1xuhao1
committed
Update from PGO.
1 parent 15e6627 commit 5b418a9

File tree

5 files changed

+129
-42
lines changed

5 files changed

+129
-42
lines changed

‎launch/taichislam-d435.launch

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,11 @@
3030
<param name="disp/max_mesh" value="$(arg max_mesh)" type="int" />
3131

3232
<remap from="~depth" to="/camera/depth/image_rect_raw" />
33-
<remap from="~pose" to="/vins_estimator/camera_pose" />
33+
<!-- <remap from="~pose" to="/vins_estimator/camera_pose" /> -->
34+
<remap from="~pose" to="/d2vins/camera_pose_1" />
35+
<remap from="~traj" to="/d2pgo/pgo_traj" />
36+
<remap from="~frame_local" to="/d2vins/frame_local" />
3437
<remap from="~image" to="/camera/infra1/image_rect_raw/compressed" />
35-
<!-- <remap from="~image" to="/camera/color/image_raw/compressed" /> -->
3638

3739
<rosparam>
3840
texture_compressed: true

‎scripts/taichislam_node.py

Lines changed: 74 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,9 @@
1616
import message_filters
1717
import cv2
1818
from transformations import *
19+
from swarm_msgs.msg import DroneTraj
20+
import struct
21+
from swarm_msgs.msg import VIOFrame
1922

2023
cur_trans = None
2124
pub = None
@@ -30,7 +33,7 @@ def __init__(self):
3033
self.enable_mesher = rospy.get_param('~enable_mesher', True)
3134
self.enable_rendering = rospy.get_param('~enable_rendering', True)
3235
self.output_map = rospy.get_param('~output_map', False)
33-
36+
self.enable_submap = rospy.get_param('~enable_submap', False)
3437

3538
if cuda:
3639
ti.init(arch=ti.cuda, device_memory_fraction=0.5, dynamic_index=True, offline_cache=True)
@@ -40,6 +43,7 @@ def __init__(self):
4043
self.disp_level = 0
4144
self.count = 0
4245
self.cur_pose = None ## Naive approach, need sync!!!
46+
self.cur_frame = None ## Naive approach, need sync!!!
4347
self.initial_mapping()
4448

4549
if self.enable_rendering:
@@ -52,21 +56,7 @@ def __init__(self):
5256
self.pub_occ = rospy.Publisher('/occ', PointCloud2, queue_size=10)
5357
self.pub_tsdf_surface = rospy.Publisher('/pub_tsdf_surface', PointCloud2, queue_size=10)
5458

55-
self.pose_sub = message_filters.Subscriber('~pose', PoseStamped)
56-
self.depth_sub = message_filters.Subscriber('~depth', Image, queue_size=10)
57-
58-
if self.texture_enabled:
59-
if self.texture_compressed:
60-
self.image_sub = message_filters.Subscriber('~image', CompressedImage, queue_size=10)
61-
else:
62-
self.image_sub = message_filters.Subscriber('~image', Image)
63-
self.ts = message_filters.ApproximateTimeSynchronizer([self.depth_sub, self.image_sub, self.pose_sub], 10, slop=0.03)
64-
self.ts.registerCallback(self.process_depth_image_pose)
65-
66-
else:
67-
self.ts = message_filters.ApproximateTimeSynchronizer([self.depth_sub, self.pose_sub], 10, slop=0.03)
68-
self.ts.registerCallback(self.process_depth_pose)
69-
59+
self.init_subscribers()
7060
fx_dep = rospy.get_param('Kdepth/fx', 384.2377014160156)
7161
fy_dep = rospy.get_param('Kdepth/fy', 384.2377014160156)
7262
cx_dep = rospy.get_param('Kdepth/cx', 323.4873046875)
@@ -81,6 +71,36 @@ def __init__(self):
8171
self.K = np.array([fx_dep, 0.0, cx_dep, 0.0, fy_dep, cy_dep, 0.0, 0.0, 1.0])
8272
self.Kcolor = np.array([fx_color, 0.0, cx_color, 0.0, fy_color, cy_color, 0.0, 0.0, 1.0])
8373

74+
def init_subscribers(self):
75+
self.depth_sub = message_filters.Subscriber('~depth', Image, queue_size=10)
76+
77+
if self.enable_submap:
78+
79+
self.frame_sub = message_filters.Subscriber('~frame_local', VIOFrame)
80+
self.traj_sub = rospy.Subscriber("~traj", DroneTraj, self.traj_callback, queue_size=10, tcp_nodelay=True)
81+
if self.texture_enabled:
82+
if self.texture_compressed:
83+
self.image_sub = message_filters.Subscriber('~image', CompressedImage, queue_size=10)
84+
else:
85+
self.image_sub = message_filters.Subscriber('~image', Image)
86+
self.ts = message_filters.ApproximateTimeSynchronizer([self.depth_sub, self.image_sub, self.frame_sub], 10, slop=0.03)
87+
self.ts.registerCallback(self.process_depth_image_frame)
88+
else:
89+
self.ts = message_filters.ApproximateTimeSynchronizer([self.depth_sub, self.frame_sub], 10, slop=0.03)
90+
self.ts.registerCallback(self.process_depth_frame)
91+
else:
92+
self.pose_sub = message_filters.Subscriber('~pose', PoseStamped)
93+
if self.texture_enabled:
94+
if self.texture_compressed:
95+
self.image_sub = message_filters.Subscriber('~image', CompressedImage, queue_size=10)
96+
else:
97+
self.image_sub = message_filters.Subscriber('~image', Image)
98+
self.ts = message_filters.ApproximateTimeSynchronizer([self.depth_sub, self.image_sub, self.pose_sub], 10, slop=0.03)
99+
self.ts.registerCallback(self.process_depth_image_pose)
100+
else:
101+
self.ts = message_filters.ApproximateTimeSynchronizer([self.depth_sub, self.pose_sub], 10, slop=0.03)
102+
self.ts.registerCallback(self.process_depth_pose)
103+
84104
def get_general_mapping_opts(self):
85105
max_disp_particles = rospy.get_param('~disp/max_disp_particles', 8000000)
86106
occupy_thres = rospy.get_param('~occupy_thres', 0)
@@ -124,7 +144,6 @@ def get_submap_opts(self):
124144
return opts
125145

126146
def initial_mapping(self):
127-
self.enable_submap = rospy.get_param('~enable_submap', False)
128147
self.mapping_type = rospy.get_param('~mapping_type', 'tsdf')
129148
self.texture_enabled = rospy.get_param('~texture_enabled', True)
130149
max_mesh = rospy.get_param('~disp/max_mesh', 1000000)
@@ -151,6 +170,7 @@ def initial_mapping(self):
151170
if self.enable_mesher:
152171
self.mesher = MarchingCubeMesher(self.mapping, max_mesh, tsdf_surface_thres=self.voxel_size*5)
153172

173+
#TODO: Move test to test.py
154174
def test_mesher(self):
155175
self.mapping.init_sphere()
156176
self.mesher.generate_mesh(1)
@@ -162,14 +182,33 @@ def update_test_mesher(self):
162182
self.mapping.cvt_TSDF_to_voxels_slice(self.render.slice_z, 100)
163183
self.render.set_particles(self.mapping.export_TSDF_xyz, self.mapping.export_color)
164184

185+
def process_depth_frame(self, depth_msg, frame):
186+
self.cur_frame = frame
187+
self.depth_msg = depth_msg
188+
self.rgb_array = np.array([], dtype=int)
189+
self.texture_image = np.array([], dtype=int)
190+
191+
def process_depth_image_frame(self, depth_msg, image, frame):
192+
if type(image) == CompressedImage:
193+
np_arr = np.frombuffer(image.data, np.uint8)
194+
self.texture_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
195+
else:
196+
np_arr = np.frombuffer(image.data, np.uint8)
197+
np_arr = np_arr.reshape((image.height, image.width, -1))
198+
self.texture_image = np_arr
199+
self.depth_msg = depth_msg
200+
self.rgb_array = np.array([], dtype=int)
201+
self.cur_frame = frame
165202

166203
def process_depth_pose(self, depth_msg, pose):
204+
#TODO: frame from pose
167205
self.cur_pose = pose
168206
self.depth_msg = depth_msg
169207
self.rgb_array = np.array([], dtype=int)
170208
self.texture_image = np.array([], dtype=int)
171209

172210
def process_depth_image_pose(self, depth_msg, image, pose):
211+
#TODO: frame from pose
173212
if type(image) == CompressedImage:
174213
np_arr = np.frombuffer(image.data, np.uint8)
175214
self.texture_image = cv2.imdecode(np_arr, cv2.IMREAD_COLOR)
@@ -190,23 +229,18 @@ def process_pcl_pose(self, msg, pose):
190229
# self.taichimapping_pcl_callback(pose, xyz_array, rgb_array)
191230

192231
def update(self):
193-
if self.cur_pose is not None:
194-
self.taichimapping_depth_callback(self.cur_pose, self.depth_msg, self.rgb_array)
232+
if self.cur_frame is not None:
233+
self.taichimapping_depth_callback(self.cur_frame, self.depth_msg, self.rgb_array)
195234
self.cur_pose = None
196235
else:
197236
if self.mapping_type == "tsdf" and self.enable_rendering:
198237
if self.render.enable_slice_z:
199238
self.mapping.cvt_TSDF_to_voxels_slice(self.render.slice_z)
200239
self.render.set_particles(self.mapping.export_TSDF_xyz, self.mapping.export_color)
201240

202-
203-
204-
def taichimapping_depth_callback(self, pose, depth_msg, rgb_array=None):
241+
def taichimapping_depth_callback(self, frame, depth_msg, rgb_array=None):
205242
mapping = self.mapping
206-
207243
start_time = time.time()
208-
209-
_R, _T = pose_msg_to_numpy(pose.pose)
210244

211245
t_pcl2npy = (time.time() - start_time)*1000
212246
start_time = time.time()
@@ -216,7 +250,15 @@ def taichimapping_depth_callback(self, pose, depth_msg, rgb_array=None):
216250
depthmap = np.frombuffer(depth_msg.data, dtype=np.uint16)
217251
depthmap = np.reshape(depthmap, (h, w))
218252

219-
mapping.recast_depth_to_map(_R, _T, depthmap, self.texture_image, w, h, self.K, self.Kcolor)
253+
if self.enable_submap:
254+
pose = pose_msg_to_numpy(frame.odom.pose.pose)
255+
frame_id = frame.frame_id
256+
print("process frame", frame_id)
257+
ext = pose_msg_to_numpy(frame.extrinsics[0])
258+
mapping.recast_depth_to_map_by_frame(frame_id, frame.is_keyframe, pose, ext, depthmap, self.texture_image, w, h, self.K, self.Kcolor)
259+
else:
260+
R, T = pose_msg_to_numpy(frame.odom.pose.pose)
261+
mapping.recast_depth_to_map_by(frame_id, R, T, depthmap, self.texture_image, w, h, self.K, self.Kcolor)
220262

221263
t_recast = (time.time() - start_time)*1000
222264

@@ -297,6 +339,12 @@ def taichimapping_pcl_callback(self, pose, xyz_array, rgb_array=None):
297339

298340
self.count += 1
299341
print(f"Time: pcl2npy {t_pcl2npy:.1f}ms t_recast {t_recast:.1f}ms ms t_v2p {t_v2p:.1f}ms t_pubros {t_pubros:.1f}ms t_render {t_render:.1f}ms")
342+
343+
def traj_callback(self, traj):
344+
frame_poses = {}
345+
for i in range(len(traj.frame_ids)):
346+
frame_poses[traj.frame_ids[i]] = pose_msg_to_numpy(traj.poses[i])
347+
self.mapping.set_frame_poses(frame_poses)
300348

301349
def pub_to_ros_surface(self, pos_, colors_, TEXTURE_ENABLED):
302350
if TEXTURE_ENABLED:

‎taichi_slam/mapping/dense_sdf.py

Lines changed: 21 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -455,31 +455,39 @@ def cvt_occupy_to_voxels(self):
455455
self.cvt_TSDF_surface_to_voxels()
456456

457457
def cvt_TSDF_surface_to_voxels(self):
458-
self.cvt_TSDF_surface_to_voxels_kernel(self.clear_last_TSDF_exporting)
458+
self.cvt_TSDF_surface_to_voxels_kernel(self.num_export_TSDF_particles,
459+
self.export_TSDF_xyz, self.export_color, self.max_disp_particles,
460+
self.clear_last_TSDF_exporting, False)
459461
self.clear_last_TSDF_exporting = False
460462

463+
def cvt_TSDF_surface_to_voxels_to(self, num_export_TSDF_particles, max_disp_particles, export_TSDF_xyz, export_color):
464+
self.cvt_TSDF_surface_to_voxels_kernel(num_export_TSDF_particles,
465+
export_TSDF_xyz, export_color, max_disp_particles, False, True)
466+
461467
@ti.kernel
462-
def cvt_TSDF_surface_to_voxels_kernel(self, clear_last:ti.template()):
468+
def cvt_TSDF_surface_to_voxels_kernel(self, num_export_TSDF_particles:ti.template(), export_TSDF_xyz:ti.template(),
469+
export_color:ti.template(), max_disp_particles:ti.template(), clear_last:ti.template(), add_to_cur:ti.template()):
463470
# Number for TSDF
464471
if clear_last:
465-
for i in range(self.num_export_TSDF_particles[None]):
466-
self.export_color[i] = ti.Vector([0.5, 0.5, 0.5], ti.f32)
467-
self.export_TSDF_xyz[i] = ti.Vector([-100000, -100000, -100000], ti.f32)
472+
for i in range(num_export_TSDF_particles[None]):
473+
export_color[i] = ti.Vector([0.5, 0.5, 0.5], ti.f32)
474+
export_TSDF_xyz[i] = ti.Vector([-100000, -100000, -100000], ti.f32)
475+
476+
if not add_to_cur:
477+
num_export_TSDF_particles[None] = 0
468478

469-
self.num_export_TSDF_particles[None] = 0
470479
for s, i, j, k in self.TSDF:
471480
if s == self.active_submap_id[None]:
472481
if self.TSDF_observed[s, i, j, k]:
473482
if ti.abs(self.TSDF[s, i, j, k] ) < self.tsdf_surface_thres:
474-
index = ti.atomic_add(self.num_export_TSDF_particles[None], 1)
475-
if self.num_export_TSDF_particles[None] < self.max_disp_particles:
476-
self.export_TSDF[index] = self.TSDF[s, i, j, k]
483+
index = ti.atomic_add(num_export_TSDF_particles[None], 1)
484+
if num_export_TSDF_particles[None] < max_disp_particles:
477485
if ti.static(self.TEXTURE_ENABLED):
478-
self.export_color[index] = self.color[s, i, j, k]
486+
export_color[index] = self.color[s, i, j, k]
479487
if ti.static(self.is_global_map):
480-
self.export_TSDF_xyz[index] = self.i_j_k_to_xyz(i, j, k)
488+
export_TSDF_xyz[index] = self.i_j_k_to_xyz(i, j, k)
481489
else:
482-
self.export_TSDF_xyz[index] = self.submap_i_j_k_to_xyz(s, i, j, k)
490+
export_TSDF_xyz[index] = self.submap_i_j_k_to_xyz(s, i, j, k)
483491

484492
@ti.kernel
485493
def cvt_ESDF_to_voxels_slice(self, z: ti.template()):
@@ -515,6 +523,7 @@ def cvt_TSDF_to_voxels_slice_kernel(self, z: ti.template(), dz:ti.template()):
515523

516524
@ti.kernel
517525
def fuse_submaps_kernel(self, TSDF:ti.template(), W_TSDF:ti.template(), TSDF_observed:ti.template(), color:ti.template()):
526+
print("submap10 R T", self.submaps_base_R[10], self.submaps_base_T[10])
518527
for s, i, j, k in TSDF:
519528
if TSDF_observed[s, i, j, k] > 0:
520529
tsdf = TSDF[s, i, j, k]

‎taichi_slam/mapping/mapping_common.py

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -80,6 +80,10 @@ def switch_to_next_submap(self):
8080
def set_base_pose_submap(self, submap_id, _R, _T):
8181
self.submaps_base_T_np[submap_id] = _T
8282
self.submaps_base_R_np[submap_id] = _R
83+
self.set_base_pose_submap_kernel(submap_id, _R, _T)
84+
85+
@ti.kernel
86+
def set_base_pose_submap_kernel(self, submap_id:ti.i16, _R:ti.types.ndarray(), _T:ti.types.ndarray()):
8387
for i in range(3):
8488
self.submaps_base_T[submap_id][i] = _T[i]
8589
for j in range(3):

‎taichi_slam/mapping/submap_mapping.py

Lines changed: 26 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,15 @@ def set_export_submap(self, new_submap):
6262
self.num_export_TSDF_particles = new_submap.num_export_TSDF_particles
6363
else:
6464
self.export_x = new_submap.export_x
65+
66+
def set_frame_poses(self, frame_poses):
67+
print("[SubmapMapping] Update frame poses from PGO")
68+
for frame_id in frame_poses:
69+
R = frame_poses[frame_id][0]
70+
T = frame_poses[frame_id][1]
71+
if frame_id in self.submaps:
72+
print(f"[SubmapMapping] Update pose frame {frame_id} submap {self.submaps[frame_id]}")
73+
self.global_map.set_base_pose_submap(self.submaps[frame_id], R, T)
6574

6675
def create_new_submap(self, frame_id, R, T):
6776
if self.first_init:
@@ -71,16 +80,18 @@ def create_new_submap(self, frame_id, R, T):
7180
self.submap_collection.clear_last_TSDF_exporting = True
7281
self.local_to_global()
7382
submap_id = self.submap_collection.get_active_submap_id()
74-
self.submap_collection.set_base_pose_submap(submap_id, R, T)
7583
self.global_map.set_base_pose_submap(submap_id, R, T)
84+
self.submap_collection.set_base_pose_submap(submap_id, R, T)
7685
self.submaps[frame_id] = submap_id
7786

7887
print(f"[SubmapMapping] Created new submap, now have {submap_id+1} submaps")
7988
return self.submap_collection
8089

81-
def need_create_new_submap(self, R, T):
90+
def need_create_new_submap(self, is_keyframe, R, T):
8291
if self.frame_count == 0:
8392
return True
93+
if not is_keyframe:
94+
return False
8495
if self.frame_count % self.keyframe_step == 0:
8596
return True
8697
return False
@@ -91,6 +102,17 @@ def recast_pcl_to_map(self, R, T, xyz_array, rgb_array, n):
91102
def local_to_global(self):
92103
self.global_map.fuse_submaps(self.submap_collection)
93104

105+
def recast_depth_to_map_by_frame(self, frame_id, is_keyframe, pose, ext, depthmap, texture, w, h, K, Kcolor):
106+
R, T = pose
107+
R_ext, T_ext = ext
108+
if self.need_create_new_submap(is_keyframe, R, T):
109+
#In early debug we use framecount as frameid
110+
self.create_new_submap(frame_id, R, T)
111+
Rcam = R @ R_ext
112+
Tcam = T + R @ T_ext
113+
self.submap_collection.recast_depth_to_map(Rcam, Tcam, depthmap, texture, w, h, K, Kcolor)
114+
self.frame_count += 1
115+
94116
def recast_depth_to_map(self, R, T, depthmap, texture, w, h, K, Kcolor):
95117
if self.need_create_new_submap(R, T):
96118
#In early debug we use framecount as frameid
@@ -109,5 +131,7 @@ def cvt_TSDF_surface_to_voxels(self):
109131
if len(self.submaps) > 0:
110132
if self.exporting_global:
111133
self.global_map.cvt_TSDF_surface_to_voxels()
134+
self.submap_collection.cvt_TSDF_surface_to_voxels_to(self.global_map.num_export_TSDF_particles,
135+
self.global_map.max_disp_particles, self.export_TSDF_xyz, self.export_color)
112136
else:
113137
self.submap_collection.cvt_TSDF_surface_to_voxels()

0 commit comments

Comments
 (0)