def rviz_kalman(demo_dir, bag_file, data_file, calib_file, freq, use_rgbd, use_smoother, use_spline, customized_shift, single_camera):
'''
For rgbd, data_file and bag_file are redundant
Otherwise, demo_dir is redundant
'''
if use_rgbd:
bag_file = osp.join(demo_dir, 'demo.bag')
rgbd1_dir = osp.join(demo_dir, 'camera_#1')
rgbd2_dir = osp.join(demo_dir, 'camera_#2')
data_file = osp.join(demo_dir, 'demo.data')
bag = rosbag.Bag(bag_file)
else:
bag = rosbag.Bag(bag_file)
dat = cp.load(open(data_file))
grippers = dat.keys()
pub = rospy.Publisher('/point_cloud1', PointCloud2)
pub2= rospy.Publisher('/point_cloud2', PointCloud2)
c1_tfm_pub = {}
c2_tfm_pub = {}
hydra_tfm_pub = {}
T_filt = {}
ang_strm = {}
ar1_strm = {}
ar2_strm = {}
hy_strm = {}
smooth_hy = {}
## publishers for unfiltered-data:
for lr in grippers:
c1_tfm_pub[lr] = rospy.Publisher('/%s_ar1_estimate'%(lr), PoseStamped)
c2_tfm_pub[lr] = rospy.Publisher('/%s_ar2_estimate'%(lr), PoseStamped)
hydra_tfm_pub[lr] = rospy.Publisher('/%s_hydra_estimate'%(lr), PoseStamped)
_, _, _, ar1_strm[lr], ar2_strm[lr], hy_strm = relative_time_streams(data_file, lr, freq, single_camera)
## run the kalman filter:
nsteps, tmin, F_means, S,A,R = run_kalman_filter(data_file, lr, freq, use_spline, True, single_camera)
## run the kalman smoother:
S_means, _ = smoother(A, R, F_means, S)
X_kf = np.array(F_means)
X_kf = np.reshape(X_kf, (X_kf.shape[0], X_kf.shape[1])).T
X_ks = np.array(S_means)
X_ks = np.reshape(X_ks, (X_ks.shape[0], X_ks.shape[1])).T
# Shifting between filter and smoother:
if customized_shift != None:
shift = customized_shift
else:
shift = correlation_shift(X_kf, X_ks)
X_ks = np.roll(X_ks,shift,axis=1)
X_ks[:,:shift] = X_ks[:,shift][:,None]
if use_smoother:
T_filt[lr] = state_to_hmat(list(X_ks.T))
else:
T_filt[lr] = state_to_hmat(list(X_kf.T))
## load the potentiometer-angle stream:
pot_data = cp.load(open(data_file))[lr]['pot_angles']
ang_ts = np.array([tt[1] for tt in pot_data]) ## time-stamps
ang_vals = [tt[0] for tt in pot_data] ## angles
# plt.plot(ang_vals)
# plt.show()
ang_vals = [0*open_frac(x) for x in ang_vals]
ang_strm[lr] = streamize(ang_vals, ang_ts, freq, lambda x : x[-1], tmin)
if use_spline:
smooth_hy[lr] = (t for t in fit_spline_to_stream(hy_strm, nsteps))
else:
smooth_hy[lr] = hy_strm
## get the point-cloud stream
cam1_frame_id = '/camera1_rgb_optical_frame'
cam2_frame_id = '/camera2_rgb_optical_frame'
if use_rgbd:
pc1_strm = streamize_rgbd_pc(rgbd1_dir, cam1_frame_id, freq, tmin)
if single_camera:
pc2_strm = streamize_rgbd_pc(None, cam2_frame_id, freq, tmin)
else:
pc2_strm = streamize_rgbd_pc(rgbd2_dir, cam2_frame_id, freq, tmin)
else:
pc1_strm = streamize_pc(bag, '/camera1/depth_registered/points', freq, tmin)
if single_camera:
pc2_strm = streamize_pc(bag, None, freq, tmin)
else:
pc2_strm = streamize_pc(bag, '/camera2/depth_registered/points', freq, tmin)
## get the relative-transforms between the cameras:
cam_tfm = get_cam_transform(calib_file)
publish_static_tfm(cam1_frame_id, cam2_frame_id, cam_tfm)
## frame of the filter estimate:
sleeper = rospy.Rate(freq)
T_far = np.eye(4)
T_far[0:3,3] = [10,10,10]
handles = []
prev_ang = {'l': 0, 'r': 0}
for i in xrange(nsteps):
#raw_input("Hit next when ready.")
print "Kalman ts: ", tmin+(0.0+i)/freq
## show the point-cloud:
found_pc = False
try:
pc = pc1_strm.next()
if pc is not None:
print "pc1 ts:", pc.header.stamp.to_sec()
pc.header.stamp = rospy.Time.now()
pub.publish(pc)
found_pc = True
else:
print "pc1 ts:",None
except StopIteration:
print "pc1 ts: finished"
#print "no more point-clouds"
pass
try:
pc2 = pc2_strm.next()
if pc2 is not None:
#print "pc2 not none"
print "pc2 ts:", pc2.header.stamp.to_sec()
pc2.header.stamp = rospy.Time.now()
pub2.publish(pc2)
found_pc = True
else:
print "pc2 ts:", None
except StopIteration:
print "pc2 ts: finished"
pass
ang_vals = []
T_filt_lr = []
for lr in grippers:
ang_val = soft_next(ang_strm[lr])
if ang_val != None:
prev_ang[lr] = ang_val
ang_val = ang_val
else:
ang_val = prev_ang[lr]
ang_vals.append(ang_val)
T_filt_lr.append(T_filt[lr][i])
handles = draw_trajectory(cam1_frame_id, T_filt_lr, color=(1,1,0,1), open_fracs=ang_vals)
# draw un-filtered estimates:
for lr in grippers:
ar1_est = soft_next(ar1_strm[lr])
if ar1_est != None:
c1_tfm_pub[lr].publish(pose_to_stamped_pose(hmat_to_pose(ar1_est), cam1_frame_id))
else:
c1_tfm_pub[lr].publish(pose_to_stamped_pose(hmat_to_pose(T_far), cam1_frame_id))
ar2_est = soft_next(ar2_strm[lr])
if ar2_est != None:
c2_tfm_pub[lr].publish(pose_to_stamped_pose(hmat_to_pose(ar2_est), cam1_frame_id))
else:
c2_tfm_pub[lr].publish(pose_to_stamped_pose(hmat_to_pose(T_far), cam1_frame_id))
hy_est = soft_next(smooth_hy[lr])
if hy_est != None:
hydra_tfm_pub[lr].publish(pose_to_stamped_pose(hmat_to_pose(hy_est), cam1_frame_id))
else:
hydra_tfm_pub[lr].publish(pose_to_stamped_pose(hmat_to_pose(T_far), cam1_frame_id))
sleeper.sleep()
demo_file = 'demo1.data' fname = osp.join(demo_dir, demo_file) dat = cp.load(open(fname)) scale = 10. ds = 3 cam1_tfms = scale_tfms([tt[0] for tt in dat['camera1']][::ds], scale) cam2_tfms = scale_tfms([tt[0] for tt in dat['camera2']][::ds], scale) hydra_tfms = scale_tfms([tt[0] for tt in dat['hydra']][::50], scale) frame = 'camera_link' handles = [] handles.extend(rv.draw_trajectory(frame, cam1_tfms, color=(1,0,0,0.5))) handles.extend(rv.draw_trajectory(frame, cam2_tfms, color=(0,1,0,0.5))) handles.extend(rv.draw_trajectory(frame, hydra_tfms, color=(0,0,1,0.5))) ## draw curves connecting the grippers: cam1_poses = rv.hmats_to_pose_array(cam1_tfms, frame) cam1_colors = [(i/float(len(cam1_poses.poses)),0,0,0.5) for i in xrange(len(cam1_poses.poses))] handles.append(rv.rviz.draw_curve(cam1_poses, rgba=cam1_colors)) cam2_poses = rv.hmats_to_pose_array(cam2_tfms, frame) cam2_colors = [(0,i/float(len(cam1_poses.poses)),0,0.5) for i in xrange(len(cam2_poses.poses))] handles.append(rv.rviz.draw_curve(cam2_poses, rgba=cam2_colors)) hydra_poses = rv.hmats_to_pose_array(hydra_tfms, frame) hydra_colors = [(0,0,i/float(len(hydra_poses.poses)),0.5) for i in xrange(len(hydra_poses.poses))] handles.append(rv.rviz.draw_curve(hydra_poses, rgba=hydra_colors))
