Skip to content

Commit

Permalink
[FIX] automatically respawn if tracking fails at runtime
Browse files Browse the repository at this point in the history
  • Loading branch information
@hollydinkel
hollydinkel committed Oct 30, 2023
1 parent 0d9ddf7 commit b77d80f
Show file tree
Hide file tree
Showing 2 changed files with 73 additions and 68 deletions.
4 changes: 2 additions & 2 deletions launch/trackdlo.launch
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
<arg name="multi_color_dlo" default="false" />

<!-- load parameters to corresponding nodes -->
<node name="trackdlo" pkg="trackdlo" type="trackdlo" output="screen">
<node name="trackdlo" pkg="trackdlo" type="trackdlo" output="screen" respawn="true">
<param name="camera_info_topic" type="string" value="$(arg camera_info_topic)" />
<param name="rgb_topic" type="string" value="$(arg rgb_topic)" />
<param name="depth_topic" type="string" value="$(arg depth_topic)" />
Expand Down Expand Up @@ -61,7 +61,7 @@
</node>

<!-- launch python node for initialization -->
<node name="init_tracker" pkg="trackdlo" type="initialize.py" output="screen">
<node name="init_tracker" pkg="trackdlo" type="initialize.py" output="screen" respawn="true">
<param name="camera_info_topic" type="string" value="$(arg camera_info_topic)" />
<param name="rgb_topic" type="string" value="$(arg rgb_topic)" />
<param name="depth_topic" type="string" value="$(arg depth_topic)" />
Expand Down
137 changes: 71 additions & 66 deletions trackdlo/src/initialize.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,6 @@
import time
import cv2
import numpy as np
import time

from visualization_msgs.msg import MarkerArray
from scipy import interpolate
Expand Down Expand Up @@ -68,75 +67,80 @@ def callback (rgb, depth):
# color thresholding
mask = color_thresholding(hsv_image, cur_depth)

start_time = time.time()
mask = cv2.cvtColor(mask.copy(), cv2.COLOR_GRAY2BGR)

# returns the pixel coord of points (in order). a list of lists
img_scale = 1
extracted_chains = extract_connected_skeleton(visualize_initialization_process, mask, img_scale=img_scale, seg_length=8, max_curvature=25)

all_pixel_coords = []
for chain in extracted_chains:
all_pixel_coords += chain
print('Finished extracting chains. Time taken:', time.time()-start_time)

all_pixel_coords = np.array(all_pixel_coords) * img_scale
all_pixel_coords = np.flip(all_pixel_coords, 1)

pc_z = cur_depth[tuple(map(tuple, all_pixel_coords.T))] / 1000.0
f = proj_matrix[0, 0]
cx = proj_matrix[0, 2]
cy = proj_matrix[1, 2]
pixel_x = all_pixel_coords[:, 1]
pixel_y = all_pixel_coords[:, 0]

pc_x = (pixel_x - cx) * pc_z / f
pc_y = (pixel_y - cy) * pc_z / f
extracted_chains_3d = np.vstack((pc_x, pc_y))
extracted_chains_3d = np.vstack((extracted_chains_3d, pc_z))
extracted_chains_3d = extracted_chains_3d.T

# do not include those without depth values
extracted_chains_3d = extracted_chains_3d[((extracted_chains_3d[:, 0] != 0) | (extracted_chains_3d[:, 1] != 0) | (extracted_chains_3d[:, 2] != 0))]

if multi_color_dlo:
depth_threshold = 0.58 # m
extracted_chains_3d = extracted_chains_3d[extracted_chains_3d[:, 2] > depth_threshold]

# tck, u = interpolate.splprep(extracted_chains_3d.T, s=0.001)
tck, u = interpolate.splprep(extracted_chains_3d.T, s=0.0005)
# 1st fit, less points
u_fine = np.linspace(0, 1, 300) # <-- num fit points
x_fine, y_fine, z_fine = interpolate.splev(u_fine, tck)
spline_pts = np.vstack((x_fine, y_fine, z_fine)).T

# 2nd fit, higher accuracy
num_true_pts = int(np.sum(np.sqrt(np.sum(np.square(np.diff(spline_pts, axis=0)), axis=1))) * 1000)
u_fine = np.linspace(0, 1, num_true_pts) # <-- num true points
x_fine, y_fine, z_fine = interpolate.splev(u_fine, tck)
spline_pts = np.vstack((x_fine, y_fine, z_fine)).T
total_spline_len = np.sum(np.sqrt(np.sum(np.square(np.diff(spline_pts, axis=0)), axis=1)))

init_nodes = spline_pts[np.linspace(0, num_true_pts-1, num_of_nodes).astype(int)]

results = ndarray2MarkerArray(init_nodes, result_frame_id, [1, 150/255, 0, 0.75], [0, 1, 0, 0.75])
results_pub.publish(results)

# add color
pc_rgba = struct.unpack('I', struct.pack('BBBB', 255, 40, 40, 255))[0]
pc_rgba_arr = np.full((len(init_nodes), 1), pc_rgba)
pc_colored = np.hstack((init_nodes, pc_rgba_arr)).astype(object)
pc_colored[:, 3] = pc_colored[:, 3].astype(int)

header.stamp = rospy.Time.now()
converted_points = pcl2.create_cloud(header, fields, pc_colored)
pc_pub.publish(converted_points)

rospy.signal_shutdown('Finished initial node set computation.')
try:
start_time = time.time()
mask = cv2.cvtColor(mask.copy(), cv2.COLOR_GRAY2BGR)

# returns the pixel coord of points (in order). a list of lists
img_scale = 1
extracted_chains = extract_connected_skeleton(visualize_initialization_process, mask, img_scale=img_scale, seg_length=8, max_curvature=25)

all_pixel_coords = []
for chain in extracted_chains:
all_pixel_coords += chain
print('Finished extracting chains. Time taken:', time.time()-start_time)

all_pixel_coords = np.array(all_pixel_coords) * img_scale
all_pixel_coords = np.flip(all_pixel_coords, 1)

pc_z = cur_depth[tuple(map(tuple, all_pixel_coords.T))] / 1000.0
f = proj_matrix[0, 0]
cx = proj_matrix[0, 2]
cy = proj_matrix[1, 2]
pixel_x = all_pixel_coords[:, 1]
pixel_y = all_pixel_coords[:, 0]

pc_x = (pixel_x - cx) * pc_z / f
pc_y = (pixel_y - cy) * pc_z / f
extracted_chains_3d = np.vstack((pc_x, pc_y))
extracted_chains_3d = np.vstack((extracted_chains_3d, pc_z))
extracted_chains_3d = extracted_chains_3d.T

# do not include those without depth values
extracted_chains_3d = extracted_chains_3d[((extracted_chains_3d[:, 0] != 0) | (extracted_chains_3d[:, 1] != 0) | (extracted_chains_3d[:, 2] != 0))]

if multi_color_dlo:
depth_threshold = 0.58 # m
extracted_chains_3d = extracted_chains_3d[extracted_chains_3d[:, 2] > depth_threshold]

# tck, u = interpolate.splprep(extracted_chains_3d.T, s=0.001)
tck, u = interpolate.splprep(extracted_chains_3d.T, s=0.0005)
# 1st fit, less points
u_fine = np.linspace(0, 1, 300) # <-- num fit points
x_fine, y_fine, z_fine = interpolate.splev(u_fine, tck)
spline_pts = np.vstack((x_fine, y_fine, z_fine)).T

# 2nd fit, higher accuracy
num_true_pts = int(np.sum(np.sqrt(np.sum(np.square(np.diff(spline_pts, axis=0)), axis=1))) * 1000)
u_fine = np.linspace(0, 1, num_true_pts) # <-- num true points
x_fine, y_fine, z_fine = interpolate.splev(u_fine, tck)
spline_pts = np.vstack((x_fine, y_fine, z_fine)).T
total_spline_len = np.sum(np.sqrt(np.sum(np.square(np.diff(spline_pts, axis=0)), axis=1)))

init_nodes = spline_pts[np.linspace(0, num_true_pts-1, num_of_nodes).astype(int)]

results = ndarray2MarkerArray(init_nodes, result_frame_id, [1, 150/255, 0, 0.75], [0, 1, 0, 0.75])
results_pub.publish(results)

# add color
pc_rgba = struct.unpack('I', struct.pack('BBBB', 255, 40, 40, 255))[0]
pc_rgba_arr = np.full((len(init_nodes), 1), pc_rgba)
pc_colored = np.hstack((init_nodes, pc_rgba_arr)).astype(object)
pc_colored[:, 3] = pc_colored[:, 3].astype(int)

header.stamp = rospy.Time.now()
converted_points = pcl2.create_cloud(header, fields, pc_colored)
pc_pub.publish(converted_points)
except:
rospy.logerr("Failed to extract splines.")
rospy.signal_shutdown('Stopping initialization.')

if __name__=='__main__':
rospy.init_node('init_tracker', anonymous=True)

global new_messages
new_messages=False

num_of_nodes = rospy.get_param('/init_tracker/num_of_nodes')
multi_color_dlo = rospy.get_param('/init_tracker/multi_color_dlo')
camera_info_topic = rospy.get_param('/init_tracker/camera_info_topic')
Expand All @@ -156,6 +160,7 @@ def callback (rgb, depth):
camera_info_sub = rospy.Subscriber(camera_info_topic, CameraInfo, camera_info_callback)
rgb_sub = message_filters.Subscriber(rgb_topic, Image)
depth_sub = message_filters.Subscriber(depth_topic, Image)
trackdlo_sub = rospy.Subscriber('/trackdlo/results_pc', PointCloud2, received_new_messages)

# header
header = std_msgs.msg.Header()
Expand Down

0 comments on commit b77d80f

Please sign in to comment.