<s> from lfd . rapprentice import ros2rave , func_utils , berkeley_pr2 <EOL> import fastrapp <EOL> import numpy as np <EOL> import cv2 <EOL> import openravepy <EOL> import os . path as osp <EOL> def extract_joints ( bag ) : <EOL> """<STR_LIT>""" <EOL> traj = [ ] <EOL> stamps = [ ] <EOL> for ( _ , msg , _ ) in bag . read_messages ( topics = [ '<STR_LIT>' ] ) : <EOL> traj . append ( msg . position ) <EOL> stamps . append ( msg . header . stamp . to_sec ( ) ) <EOL> assert len ( traj ) > <NUM_LIT:0> <EOL> names = msg . name <EOL> return names , stamps , traj <EOL> def extract_joy ( bag ) : <EOL> """<STR_LIT>""" <EOL> stamps = [ ] <EOL> meanings = [ ] <EOL> button2meaning = { <EOL> <NUM_LIT:12> : "<STR_LIT>" , <EOL> <NUM_LIT:0> : "<STR_LIT:start>" , <EOL> <NUM_LIT:3> : "<STR_LIT>" , <EOL> <NUM_LIT:7> : "<STR_LIT>" , <EOL> <NUM_LIT:5> : "<STR_LIT>" , <EOL> <NUM_LIT:15> : "<STR_LIT>" , <EOL> <NUM_LIT> : "<STR_LIT>" , <EOL> <NUM_LIT> : "<STR_LIT>" <EOL> } <EOL> check_buttons = button2meaning . keys ( ) <EOL> message_stream = bag . read_messages ( topics = [ '<STR_LIT>' ] ) <EOL> ( _ , last_msg , _ ) = message_stream . next ( ) <EOL> for ( _ , msg , _ ) in message_stream : <EOL> for i in check_buttons : <EOL> if msg . buttons [ i ] and not last_msg . buttons [ i ] : <EOL> stamps . append ( msg . header . stamp . to_sec ( ) ) <EOL> meanings . append ( button2meaning [ i ] ) <EOL> last_msg = msg <EOL> return stamps , meanings <EOL> def find_disjoint_subsequences ( li , seq ) : <EOL> """<STR_LIT>""" <EOL> subseqs = [ ] <EOL> cur_subseq_inds = [ ] <EOL> for ( i_el , el ) in enumerate ( li ) : <EOL> if el == seq [ len ( cur_subseq_inds ) ] : <EOL> cur_subseq_inds . append ( i_el ) <EOL> if len ( cur_subseq_inds ) == len ( seq ) : <EOL> subseqs . append ( cur_subseq_inds ) <EOL> cur_subseq_inds = [ ] <EOL> return subseqs <EOL> def joy_to_annotations ( stamps , meanings ) : <EOL> """<STR_LIT>""" <EOL> out = [ ] <EOL> ind_tuples = find_disjoint_subsequences ( meanings , [ "<STR_LIT>" , "<STR_LIT:start>" , "<STR_LIT>" ] ) <EOL> for tup in ind_tuples : <EOL> out . append ( { "<STR_LIT>" : stamps [ tup [ <NUM_LIT:0> ] ] , "<STR_LIT:start>" : stamps [ tup [ <NUM_LIT:1> ] ] , "<STR_LIT>" : stamps [ tup [ <NUM_LIT:2> ] ] } ) <EOL> return out <EOL> def add_kinematics_to_group ( group , linknames , manipnames , jointnames , robot ) : <EOL> "<STR_LIT>" <EOL> if robot is None : robot = get_robot ( ) <EOL> r2r = ros2rave . RosToRave ( robot , group [ "<STR_LIT>" ] [ "<STR_LIT:name>" ] ) <EOL> link2hmats = dict ( [ ( linkname , [ ] ) for linkname in linknames ] ) <EOL> links = [ robot . GetLink ( linkname ) for linkname in linknames ] <EOL> rave_traj = [ ] <EOL> rave_inds = r2r . rave_inds <EOL> for ros_vals in group [ "<STR_LIT>" ] [ "<STR_LIT>" ] : <EOL> r2r . set_values ( robot , ros_vals ) <EOL> rave_vals = r2r . convert ( ros_vals ) <EOL> robot . SetDOFValues ( rave_vals , rave_inds ) <EOL> rave_traj . append ( rave_vals ) <EOL> for ( linkname , link ) in zip ( linknames , links ) : <EOL> link2hmats [ linkname ] . append ( link . GetTransform ( ) ) <EOL> for ( linkname , hmats ) in link2hmats . items ( ) : <EOL> group . create_group ( linkname ) <EOL> group [ linkname ] [ "<STR_LIT>" ] = np . array ( hmats ) <EOL> rave_traj = np . array ( rave_traj ) <EOL> rave_ind_list = list ( rave_inds ) <EOL> for manipname in manipnames : <EOL> arm_inds = robot . GetManipulator ( manipname ) . GetArmIndices ( ) <EOL> group [ manipname ] = rave_traj [ : , [ rave_ind_list . index ( i ) for i in arm_inds ] ] <EOL> for jointname in jointnames : <EOL> joint_ind = robot . GetJointIndex ( jointname ) <EOL> group [ jointname ] = rave_traj [ : , rave_ind_list . index ( joint_ind ) ] <EOL> @ func_utils . once <EOL> def get_robot ( ) : <EOL> env = openravepy . Environment ( ) <EOL> env . Load ( "<STR_LIT>" ) <EOL> robot = env . GetRobots ( ) [ <NUM_LIT:0> ] <EOL> return robot <EOL> def add_bag_to_hdf ( bag , annotations , hdfroot , demo_name ) : <EOL> joint_names , stamps , traj = extract_joints ( bag ) <EOL> traj = np . asarray ( traj ) <EOL> stamps = np . asarray ( stamps ) <EOL> robot = get_robot ( ) <EOL> for seg_info in annotations : <EOL> group = hdfroot . create_group ( demo_name + "<STR_LIT:_>" + seg_info [ "<STR_LIT:name>" ] ) <EOL> start = seg_info [ "<STR_LIT:start>" ] <EOL> stop = seg_info [ "<STR_LIT>" ] <EOL> [ i_start , i_stop ] = np . searchsorted ( stamps , [ start , stop ] ) <EOL> stamps_seg = stamps [ i_start : i_stop + <NUM_LIT:1> ] <EOL> traj_seg = traj [ i_start : i_stop + <NUM_LIT:1> ] <EOL> sample_inds = fastrapp . resample ( traj_seg , np . arange ( len ( traj_seg ) ) , <NUM_LIT> , np . inf , np . inf ) <EOL> print "<STR_LIT>" , len ( sample_inds ) , len ( traj_seg ) <EOL> traj_ds = traj_seg [ sample_inds , : ] <EOL> stamps_ds = stamps_seg [ sample_inds ] <EOL> group [ "<STR_LIT:description>" ] = seg_info [ "<STR_LIT:description>" ] <EOL> group [ "<STR_LIT>" ] = stamps_ds <EOL> group . create_group ( "<STR_LIT>" ) <EOL> group [ "<STR_LIT>" ] [ "<STR_LIT:name>" ] = joint_names <EOL> group [ "<STR_LIT>" ] [ "<STR_LIT>" ] = traj_ds <EOL> link_names = [ "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" , "<STR_LIT>" ] <EOL> special_joint_names = [ "<STR_LIT>" , "<STR_LIT>" ] <EOL> manip_names = [ "<STR_LIT>" , "<STR_LIT>" ] <EOL> add_kinematics_to_group ( group , link_names , manip_names , special_joint_names , robot ) <EOL> def get_video_frames ( video_dir , frame_stamps ) : <EOL> video_stamps = np . loadtxt ( osp . join ( video_dir , "<STR_LIT>" ) ) <EOL> frame_inds = np . searchsorted ( video_stamps , frame_stamps ) <EOL> rgbs = [ ] <EOL> depths = [ ] <EOL> for frame_ind in frame_inds : <EOL> rgb = cv2 . imread ( osp . join ( video_dir , "<STR_LIT>" % frame_ind ) ) <EOL> assert rgb is not None <EOL> rgbs . append ( rgb ) <EOL> depth = cv2 . imread ( osp . join ( video_dir , "<STR_LIT>" % frame_ind ) , <NUM_LIT:2> ) <EOL> assert depth is not None <EOL> depths . append ( depth ) <EOL> return rgbs , depths <EOL> def add_rgbd_to_hdf ( video_dir , annotations , hdfroot , demo_name ) : <EOL> frame_stamps = [ seg_info [ "<STR_LIT>" ] for seg_info in annotations ] <EOL> rgb_imgs , depth_imgs = get_video_frames ( video_dir , frame_stamps ) <EOL> for ( i_seg , seg_info ) in enumerate ( annotations ) : <EOL> group = hdfroot [ demo_name + "<STR_LIT:_>" + seg_info [ "<STR_LIT:name>" ] ] <EOL> group [ "<STR_LIT>" ] = rgb_imgs [ i_seg ] <EOL> group [ "<STR_LIT>" ] = depth_imgs [ i_seg ] <EOL> robot = get_robot ( ) <EOL> r2r = ros2rave . RosToRave ( robot , group [ "<STR_LIT>" ] [ "<STR_LIT:name>" ] ) <EOL> r2r . set_values ( robot , group [ "<STR_LIT>" ] [ "<STR_LIT>" ] [ <NUM_LIT:0> ] ) <EOL> T_w_h = robot . GetLink ( "<STR_LIT>" ) . GetTransform ( ) <EOL> T_w_k = T_w_h . dot ( berkeley_pr2 . T_h_k ) <EOL> group [ "<STR_LIT>" ] = T_w_k </s>
