def view_hydra_demo_on_rviz (demo_type, demo_name, freq, speed, prompt=False, verbose=False):
    """
    Uses hydra_only.data for the segment to quickly visualize the demo.
    @demo_type, @demo_name: demo identification.
    @freq: basically measure of fine-ness of timesteps.
    @speed: how fast to replay demo.
    @prompt: does the user hit enter after each time step?
    """
    demo_dir = osp.join(demo_files_dir, demo_type, demo_name)
    bag_file = osp.join(demo_dir, demo_names.bag_name)
    data_file = osp.join(demo_dir, demo_names.hydra_data_name)
    calib_file = osp.join(demo_dir, demo_names.calib_name)
    with open(osp.join(demo_dir, demo_names.camera_types_name),'r') as fh: cam_types = yaml.load(fh)
    
    
    if not osp.isfile(data_file):
        yellowprint("%s does not exist for this demo. Extracting now."%demo_names.hydra_data_name)
        ed.save_hydra_only(demo_type, demo_name)
        
    with open(data_file, 'r') as fh: dat = cp.load(fh)
    
    # get grippers used
    grippers = [key for key in dat.keys() if key in 'lr']

    # data 
    rgbd_dirs = {cam:osp.join(demo_dir,demo_names.video_dir%cam) for cam in cam_types if cam_types[cam] == 'rgbd'}
    cam_frames = {cam:'/camera%i_rgb_optical_frame'%cam for cam in rgbd_dirs}
    
    tfm_pubs = {}
    hydra_dat = {}
    pot_dat = {}
    
    _, hydra_dat['l'], pot_dat['l'] = load_data(data_file, 'l', freq, speed, hydra_only=True)
    _, hydra_dat['r'], pot_dat['r'] = load_data(data_file, 'r', freq, speed, hydra_only=True)
    tmin, _, nsteps = relative_time_streams(hydra_dat.values() + pot_dat.values(), freq, speed)

    if rospy.get_name() == "/unnamed":
        rospy.init_node("visualize_demo")


    ## publishers for unfiltered-data:
    for lr in grippers:
        tfm_pubs[lr] = rospy.Publisher('/%s_hydra_estimate'%(lr), PoseStamped)

    ## get the point-cloud stream
    pc_strms = {cam:streamize_rgbd_pc(rgbd_dirs[cam], cam_frames[cam], freq, tstart=tmin,speed=speed,verbose=verbose) for cam in rgbd_dirs}
    pc_pubs = {cam:rospy.Publisher('/point_cloud%i'%cam, PointCloud2) for cam in rgbd_dirs}

    cam_tfms  = get_cam_transforms (calib_file, len(cam_types))
    for cam in rgbd_dirs:
        if cam != 1:
            publish_static_tfm(cam_frames[1], cam_frames[cam], cam_tfms[cam])

    sleeper = rospy.Rate(freq)
    T_far = np.eye(4)
    T_far[0:3,3] = [10,10,10]        
    
    handles = []
    
    dat_snext = {lr:{} for lr in grippers}
    for lr in grippers:
        dat_snext[lr]['h'] = stream_soft_next(hydra_dat[lr])
        dat_snext[lr]['pot'] = stream_soft_next(pot_dat[lr])
    
    prev_ang = {'l': 0, 'r': 0}
    for i in xrange(nsteps):
        if prompt:
            raw_input("Hit enter when ready.")
        if verbose:
            print "Time stamp: ", tmin+(0.0+i*speed)/freq
        
        ## show the point-cloud:
        for cam in pc_strms:
            try:
                pc = pc_strms[cam].next()
                if pc is not None:
                    if verbose:
                        print "pc%i ts:"%cam, pc.header.stamp.to_sec()
                    pc.header.stamp = rospy.Time.now()
                    pc_pubs[cam].publish(pc)
                else:
                    if verbose:
                        print "pc%i ts:"%cam,None
            except StopIteration:
                pass

        ests = {}
        tfms = []
        ang_vals  = []

        for lr in grippers:
            ests[lr] = dat_snext[lr]['h']()

            ang_val = dat_snext[lr]['pot']()
            if ang_val != None and not np.isnan(ang_val):
                prev_ang[lr] = ang_val
                ang_val  = ang_val
            else:
                ang_val = prev_ang[lr]
            ang_val *= 2

            if ests[lr] is None:
                tfms.append(T_far)
            else:
                tfms.append(ests[lr])
            ang_vals.append(rad_angle(ang_val))

        handles = draw_trajectory(cam_frames[1], tfms, color=(1,1,0,1), open_fracs=ang_vals)

        for lr in grippers:
            if ests[lr] is not None:
                tfm_pubs[lr].publish(conversions.pose_to_stamped_pose(conversions.hmat_to_pose(ests[lr]), cam_frames[1]))
            
        sleeper.sleep()
    
    empty_cloud = PointCloud2()
    for cam in pc_pubs: pc_pubs[cam].publish(empty_cloud)
def view_tracking_on_rviz(demo_type, demo_name, tps_model_fname, freq=30.0, speed=1.0, use_smoother=True, prompt=False, verbose=False):
    """
    Visualizes demo after kalman tracking/smoothing on rviz.
    @demo_type, @demo_name: demo identification.
    @freq: basically measure of fine-ness of timesteps.
    @speed: how fast to replay demo.
    @main: which sensor to display the marker for
    @prompt: does the user hit enter after each time step?
    """
    
    demo_dir = osp.join(demo_files_dir, demo_type, demo_name)
    bag_file = osp.join(demo_dir, demo_names.bag_name)
    traj_file = osp.join(demo_dir, demo_names.traj_name)
    calib_file = osp.join(demo_dir, demo_names.calib_name)
    with open(osp.join(demo_dir, demo_names.camera_types_name),'r') as fh: cam_types = yaml.load(fh)
    
        
    if not osp.isfile(traj_file):
        yellowprint("%s does not exist for this demo. Running kalman filter/smoother now with default args."%demo_names.traj_name)
        data_file = osp.join(demo_dir, demo_names.data_name)
        if not osp.isfile(data_file):
            yellowprint("%s does not exist for this demo. Extracting now."%demo_names.data_name)
            ed.save_observations_rgbd(demo_type, demo_name)
        filter_traj(demo_dir, tps_model_fname=tps_model_fname, save_tps=True, do_smooth=True, plot='', block=False)
    with open(traj_file, 'r') as fh: traj = cp.load(fh)
    
    # get grippers used
    grippers = traj.keys()

    if rospy.get_name() == "/unnamed":
        rospy.init_node("visualize_demo")

    # data 
    rgbd_dirs = {cam:osp.join(demo_dir,demo_names.video_dir%cam) for cam in cam_types if cam_types[cam] == 'rgbd'}
    pc_pubs = {cam:rospy.Publisher('/point_cloud%i'%cam, PointCloud2) for cam in rgbd_dirs}
    cam_frames = {cam:'/camera%i_rgb_optical_frame'%cam for cam in rgbd_dirs}
        
    cam_tfms  = get_cam_transforms (calib_file, len(cam_types))
    for cam in rgbd_dirs:
        if cam != 1:
            publish_static_tfm(cam_frames[1], cam_frames[cam], cam_tfms[cam])

    # Remove segment "done", it is just a single frame
    segs = sorted(traj[grippers[0]].keys())
    if 'done' in segs: segs.remove('done')
    
    sleeper = rospy.Rate(freq)
    T_far = np.eye(4)
    T_far[0:3,3] = [10,10,10]
    
    for seg in segs:
        if prompt:
            raw_input("Press enter for segment %s."%seg)
        else:
            yellowprint("Segment %s beginning."%seg)
            time.sleep(1)
        
        # Initializing data streams:
        traj_strms = {}
        pot_strms = {}
        tfms_key = 'tfms_s' if use_smoother else 'tfms'

        for lr in grippers:
            traj_strms[lr] = streamize(traj[lr][seg][tfms_key], traj[lr][seg]['stamps'], freq, avg_transform, speed=speed)
            # HACK
            pot_strms[lr] = streamize(traj[lr][seg]['pot_angles'][:len(traj[lr][seg]['stamps'])], traj[lr][seg]['stamps'], freq, np.mean, speed=speed)
        
        tmin, tmax, nsteps = relative_time_streams(traj_strms.values() + pot_strms.values(), freq, speed)
        


        pc_strms = {cam:streamize_rgbd_pc(rgbd_dirs[cam], cam_frames[cam], freq, tstart=tmin, tend=tmax,speed=speed,verbose=verbose) for cam in rgbd_dirs}
    
        prev_ang = {'l': 0, 'r': 0}

        dat_snext = {lr:{} for lr in grippers}
        for lr in grippers:
            dat_snext[lr]['traj'] = stream_soft_next(traj_strms[lr])
            dat_snext[lr]['pot'] = stream_soft_next(pot_strms[lr])
        

        for i in xrange(nsteps):
            if prompt:
                raw_input("Hit enter when ready.")
            if verbose:
                print "Time stamp: ", tmin+(0.0+i*speed)/freq
            
            ## show the point-cloud:
            found_pc = False
            for cam in pc_strms:
                try:
                    pc = pc_strms[cam].next()
                    if pc is not None:
                        if verbose:
                            print "pc%i ts:"%cam, pc.header.stamp.to_sec()
                        pc.header.stamp = rospy.Time.now()
                        pc_pubs[cam].publish(pc)
                        found_pc = True
                    else:
                        if verbose:
                            print "pc%i ts:"%cam,None
                except StopIteration:
                    pass
    
            tfms = []
            ang_vals  = []
    
            for lr in grippers:
		if lr == 'r':
		    continue
                tfm = dat_snext[lr]['traj']()
                
                ang_val = dat_snext[lr]['pot']()
                if ang_val != None and not np.isnan(ang_val):
                    prev_ang[lr] = ang_val
                    ang_val  = ang_val
                else:
                    ang_val = prev_ang[lr]
                ang_val *= 2

                if tfm is None:
                    tfms.append(T_far)
                else:
                    tfms.append(tfm)
                    
                ang_vals.append(rad_angle(ang_val))
    
    
            handles = draw_trajectory(cam_frames[cam], tfms, color=(1,0,0,1), open_fracs=ang_vals)
            
            time.sleep(1.0/freq)
Ejemplo n.º 3
0
def filter_traj(demo_dir, tps_model_fname, save_tps, do_smooth, plot, block):
    """
    Runs the kalman filter for BOTH the grippers and writes the demo.traj file.
    TPS_MODEL_FNAME : The name of the file to load the tps-model from
    """
    # Temp file to show that kalman filter/smoother is already being run on demo
    with open(osp.join(demo_dir, demo_names.run_kalman_temp),'w') as fh: fh.write('Running Kalman filter/smoother..')
    
    freq = 30.0
    
    rec_data, time_shifts, n_segs = prepare_kf_data(demo_dir,
                                                    freq=freq,
                                                    rem_outliers=True,
                                                    tps_correct=True, 
                                                    tps_model_fname=tps_model_fname,
                                                    plot=plot,
                                                    block=block)


    _, cam_covars, hydra_covar =  initialize_covariances(freq, demo_dir)

    KFs  = initialize_KFs(rec_data, freq)
    traj = {'l' : None, 'r':None}

    for lr in 'lr':
        lr_trajs = {}
        for iseg in xrange(n_segs):
            KF = KFs[lr][iseg]
            hydra_strm  = rec_data[lr][iseg]['hydra_strm']
            pot_strm = rec_data[lr][iseg]['pot_strm']
            cam_dat  = rec_data[lr][iseg]['cam_dat']
            nsteps   = rec_data[lr][iseg]['nsteps']


            ts, xs_kf, covars_kf, xs_smthr, covars_smthr = run_KF(KF, nsteps, freq,
                                                                  hydra_strm, cam_dat,
                                                                  hydra_covar, cam_covars,
                                                                  do_smooth,
                                                                  plot, plot_title='seg %d : %s'%(iseg, {'l':'left', 'r':'right'}[lr]),
                                                                  block=block)

            for i in range(len(ts)):
                ts[i] -= time_shifts['camera1']
                                
            Ts_kf, Ts_smthr = state_to_hmat(xs_kf), state_to_hmat(xs_smthr)
            
            pot_angles = []
            pot_ss = stream_soft_next(pot_strm)
            for _ in xrange(nsteps+1):
                pot_angles.append(pot_ss())             

            
            '''''
            Dirty hack!!!!!!!!!!!!!!!!!!!!!!!!!
            '''''
            n_pot_angles = len(pot_angles)
            if pot_angles[0] == None or np.isnan(pot_angles[0]):
                # find the first non None
                first_non_none_id = -1
                for i in xrange(n_pot_angles):
                    if pot_angles[i] != None and not np.isnan(pot_angles[i]):
                        first_non_none_id = i
                        break
                    
                if first_non_none_id != -1:
                    for i in xrange(first_non_none_id):
                        pot_angles[i] = pot_angles[first_non_none_id]
                else:
                    for i in xrange(n_pot_angles):
                        pot_angles[i] = 0
                    
            if pot_angles[-1] == None or np.isnan(pot_angles[-1]):
                first_non_none_id = -1
                for i in xrange(n_pot_angles - 1, -1, -1):
                    if pot_angles[i] != None and not np.isnan(pot_angles[i]):
                        last_non_none_id = i
                        break
                    
                if first_non_none_id != -1:
                    for i in xrange(last_non_none_id+1, n_pot_angles):
                        pot_angles[i] = pot_angles[last_non_none_id]
                else:
                    for i in xrange(n_pot_angles):
                        pot_angles[i] = 0
                    
        
            # then linear interpolation between non-None elements
            #print n_pot_angles
            #print len(ts)
            #print nsteps
            i = 0
            while i < n_pot_angles:
                if pot_angles[i] == None or np.isnan(pot_angles[i]):
                    non_none_id_0 = i - 1
                    
                    for j in xrange(i+1, n_pot_angles):
                        if pot_angles[j] != None and not np.isnan(pot_angles[j]):
                            non_none_id_1 = j
                            break
                    
                    delta = (pot_angles[non_none_id_1] - pot_angles[non_none_id_0]) / (non_none_id_1 - non_none_id_0)   
                    for j in xrange(non_none_id_0 +1, non_none_id_1):
                        pot_angles[j] = (j - non_none_id_0) * delta + pot_angles[non_none_id_0]
                        
                    #print pot_angles[non_none_id_0: non_none_id_1+1]
                    
                    i = non_none_id_1 + 1
                else:
                    i += 1
            '''
            Finish dirty hack.
            '''

            seg_traj = {"stamps"  : ts,
                        "tfms"    : Ts_kf,
                        "covars"  : covars_kf,
                        "tfms_s"  : Ts_smthr,
                        "covars_s": covars_smthr,
                        "pot_angles": pot_angles}
            
            
            seg_name = rec_data[lr][iseg]["name"]
            lr_trajs[seg_name] = seg_traj
            
        traj[lr] = lr_trajs

    traj_fname = osp.join(demo_dir, demo_names.traj_name)
    with open(traj_fname, 'wb') as f:
        cp.dump(traj, f)
    yellowprint("Saved %s."%demo_names.traj_name)
    
    os.remove(osp.join(demo_dir, demo_names.run_kalman_temp))
def view_demo_on_rviz(demo_type, demo_name, freq, speed=1.0, main='h', prompt=False, verbose=False):
    """
    Visualizes recorded demo on rviz (without kalman filter/smoother data).
    @demo_type, @demo_name: demo identification.
    @freq: basically measure of fine-ness of timesteps.
    @speed: how fast to replay demo.
    @main: which sensor to display the marker for
    @prompt: does the user hit enter after each time step?
    """
    demo_dir = osp.join(demo_files_dir, demo_type, demo_name)
    bag_file = osp.join(demo_dir, demo_names.bag_name)
    data_file = osp.join(demo_dir, demo_names.data_name)
    calib_file = osp.join(demo_dir, demo_names.calib_name)
    with open(osp.join(demo_dir, demo_names.camera_types_name),'r') as fh: cam_types = yaml.load(fh)
    
    if not osp.isfile(data_file):
        yellowprint("%s does not exist for this demo. Extracting now."%demo_names.data_name)
        ed.save_observations_rgbd(demo_type, demo_name)
    with open(data_file, 'r') as fh: dat = cp.load(fh)
    
    # get grippers used
    grippers = [key for key in dat.keys() if key in 'lr']

    # data 
    rgbd_dirs = {cam:osp.join(demo_dir,demo_names.video_dir%cam) for cam in cam_types if cam_types[cam] == 'rgbd'}
    cam_frames = {cam:'/camera%i_rgb_optical_frame'%cam for cam in rgbd_dirs}
    
    tfm_pubs = {}

    cam_dat = {}
    hydra_dat = {}
    pot_dat = {}
    
    _, cam_dat['l'], hydra_dat['l'], pot_dat['l'] = load_data(data_file, 'l', freq, speed)
    _,  cam_dat['r'], hydra_dat['r'], pot_dat['r'] = load_data(data_file, 'r', freq, speed)

    all_cam_strms = []
    for lr in 'lr':
        for cam in cam_dat[lr].keys():
            all_cam_strms.append(cam_dat[lr][cam]['stream'])
    tmin, _, nsteps = relative_time_streams(hydra_dat.values() + pot_dat.values() + all_cam_strms, freq, speed)

    if rospy.get_name() == "/unnamed":
        rospy.init_node("visualize_demo")


    ## publishers for unfiltered-data:
    for lr in grippers:
        tfm_pubs[lr] = {}
        for cam in cam_types:
            tfm_pubs[lr][cam] = rospy.Publisher('/%s_ar%i_estimate'%(lr,cam), PoseStamped)
        tfm_pubs[lr]['h'] = rospy.Publisher('/%s_hydra_estimate'%(lr), PoseStamped)

    ## get the point-cloud stream
    pc_strms = {cam:streamize_rgbd_pc(rgbd_dirs[cam], cam_frames[cam], freq, tstart=tmin,speed=speed,verbose=verbose) for cam in rgbd_dirs}
    pc_pubs = {cam:rospy.Publisher('/point_cloud%i'%cam, PointCloud2) for cam in rgbd_dirs}

#     import IPython
#     IPython.embed()

    cam_tfms  = get_cam_transforms (calib_file, len(cam_types))
    for cam in rgbd_dirs:
        if cam != 1:
            publish_static_tfm(cam_frames[1], cam_frames[cam], cam_tfms[cam])

    sleeper = rospy.Rate(freq)
    T_far = np.eye(4)
    T_far[0:3,3] = [10,10,10]        
    
    handles = []
    
    prev_ang = {'l': 0, 'r': 0}
    
    
    dat_snext = {lr:{} for lr in grippers}
    for lr in grippers:
        dat_snext[lr]['h'] = stream_soft_next(hydra_dat[lr])
        dat_snext[lr]['pot'] = stream_soft_next(pot_dat[lr])
        
        for cam in cam_types:
            dat_snext[lr][cam] = stream_soft_next(cam_dat[lr][cam]['stream'])
        
    
    
    for i in xrange(nsteps):
        if prompt:
            raw_input("Hit enter when ready.")
        if verbose:
            print "Time stamp: ", tmin+(0.0+i*speed)/freq
        
        ## show the point-cloud:
        found_pc = False
        for cam in pc_strms:
            try:
                pc = pc_strms[cam].next()
                if pc is not None:
                    if verbose:
                        print "pc%i ts:"%cam, pc.header.stamp.to_sec()
                    pc.header.stamp = rospy.Time.now()
                    pc_pubs[cam].publish(pc)
                    found_pc = True
                else:
                    if verbose:
                        print "pc%i ts:"%cam,None
            except StopIteration:
                pass

        next_est = {lr:{} for lr in grippers}
        tfms = []
        ang_vals  = []
        
        if main != 'h': main = int(main)

        for lr in grippers:
            next_est[lr]['h'] = dat_snext[lr]['h']()
            for cam in cam_types:
                next_est[lr][cam] = dat_snext[lr][cam]()

            ang_val = dat_snext[lr]['pot']()
            if ang_val != None and not np.isnan(ang_val):
                prev_ang[lr] = ang_val
                ang_val  = ang_val
            else:
                ang_val = prev_ang[lr]
            
            ang_val *= 2
            
            tfm = next_est[lr][main]
            if tfm is None:
                tfms.append(T_far)
            else:
                tfms.append(tfm)
            ang_vals.append(rad_angle(ang_val))

        handles = draw_trajectory(cam_frames[1], tfms, color=(1,1,0,1), open_fracs=ang_vals)

        for lr in grippers:
            for m,est in next_est[lr].items():
                if est != None:
                    tfm_pubs[lr][m].publish(conversions.pose_to_stamped_pose(conversions.hmat_to_pose(est), cam_frames[1]))
                else:
                    tfm_pubs[lr][m].publish(conversions.pose_to_stamped_pose(conversions.hmat_to_pose(T_far), cam_frames[1]))
        
        sleeper.sleep()
Ejemplo n.º 5
0
def run_KF(KF, nsteps, freq, hydra_strm, cam_dat, hydra_covar, cam_covars, do_smooth, plot, plot_title, block):
    """
    Runs the Kalman filter/smoother for NSTEPS using
    {hydra/rgb/rgbd}_covar as the measurement covariances.
    
    HYDRA_STRM   : TF stream from hydra
    CAM_DAT   : Dictionary of camera streams 
    do_smooth: default true
    plot: default true
    """ 
    
    dt = 1./freq
    
    ## place holders for kalman filter's output:
    xs_kf, covars_kf, ts_kf = [KF.x_filt],[KF.S_filt],[KF.t_filt]

    hydra_snext  = stream_soft_next(hydra_strm)

    cam_strms = [cinfo['stream'] for cinfo in cam_dat.values()]
    cam_names = cam_dat.keys()
    cam_types = [cinfo['type'] for cinfo in cam_dat.values()]
    cam_snext = [stream_soft_next(cstrm) for cstrm in cam_strms]
    
    for i in xrange(nsteps):
        KF.register_tf_observation(hydra_snext(), hydra_covar, do_control_update=True)

        for i in xrange(len(cam_names)):
            cam_covar = cam_covars[cam_names[i]]
            KF.register_tf_observation(cam_snext[i](), cam_covar, do_control_update=False)

        xs_kf.append(KF.x_filt)
        covars_kf.append(KF.S_filt)
        ts_kf.append(KF.t_filt)
        
    hydra_strm.reset()
    for cstrm in cam_strms:
        cstrm.reset()
    if do_smooth:

        '''
        # UNCOMMENT BELOW TO RUN KALMAN SMOOTHER:
        A,_  = KF.get_motion_mats(dt)
        R    = get_smoother_motion_covariance(freq)
        xs_smthr, covars_smthr = smoother(A, R, xs_kf, covars_kf)
        '''
        ### do low-pass filtering for smoothing:
        xs_kf_xyz = np.array(xs_kf)[:,0:3] 
        xs_kf_rpy = np.squeeze(np.array(xs_kf)[:,6:9])

        xs_lp_xyz = low_pass(xs_kf_xyz, freq)
        xs_lp_rpy = low_pass(unbreak(xs_kf_rpy), freq)

        xs_smthr  = np.c_[xs_lp_xyz, np.squeeze(np.array(xs_kf))[:,3:6], xs_lp_rpy, np.squeeze(np.array(xs_kf))[:,9:12]].tolist()
        covars_smthr = None

        if 's' in plot:
            kf_strm   = streamize(state_to_hmat(xs_kf), np.array(ts_kf), 1./hydra_strm.dt, hydra_strm.favg)
            sm_strm   = streamize(state_to_hmat(xs_smthr), np.array(ts_kf), 1./hydra_strm.dt, hydra_strm.favg)
            

            plot_tf_streams([kf_strm, sm_strm, hydra_strm]+cam_strms, ['kf', 'smoother', 'hydra']+cam_dat.keys(), styles=['-','-','.']+['.' for i in xrange(len(cam_strms))], title=plot_title, block=block)
            #plot_tf_streams([hydra_strm]+cam_strms, ['hydra']+cam_dat.keys(), block=block)

        return (ts_kf, xs_kf, covars_kf, xs_smthr, covars_smthr)
    else:
        if 's' in plot:
            kf_strm   = streamize(state_to_hmat(xs_kf), np.array(ts_kf), 1./hydra_strm.dt, hydra_strm.favg)
            plot_tf_streams([kf_strm, hydra_strm]+cam_strms, ['kf', 'hydra']+cam_dat.keys(), styles=['-','-','.']+['.' for i in xrange(len(cam_strms))], title=plot_title, block=block)
            #plot_tf_streams([hydra_strm]+cam_strms, ['hydra']+cam_dat.keys(), block=block)
        return (ts_kf, xs_kf, covars_kf, None, None)