Exemplo n.º 1
0
                                     child_frame, parent_frame)
            sleeper.sleep()
    thread.start_new_thread(spin_pub, ())



def open_frac(th):
    thmax = 33
    return th/thmax;
            
if __name__ == '__main__':
    demo_num = 6
    freq     = 30.

    data_dir = os.getenv('HD_DATA_DIR') 
    #if data_dir is None:
    bag = rosbag.Bag(hd_path + '/hd_data/demos/recorded/demo'+str(demo_num)+'.bag')
    #else:
    #	bag = rosbag.Bag(osp.join(data_dir,'demos/recorded/demo'+str(demo_num)+'.bag'))
    rospy.init_node('process_pc')
    
    ## get the point-cloud stream
    pc1_strm = streamize_pc(bag, '/camera1/depth_registered/points', freq)
    pc2_strm = streamize_pc(bag, '/camera2/depth_registered/points', freq)

    pc = None
    while pc is None:
        pc = pc1_strm.next()
    print type(pc)
    
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()