teleoperate_wii.py

#!/usr/bin/env python

"""
Human Robot Interation, Design Project II
"""
import argparse
import sys

import rospy
import tf2_ros
import math
import numpy as np

import baxter_interface
from baxter_interface import CHECK_VERSION
from wiimote.msg import State

BASE_FRAME = 'openni_depth_frame'
FRAMES = [
        'torso',
        'left_shoulder',
        'left_elbow',
        'left_hand',
        'right_shoulder',
        'right_elbow',
        'right_hand',
        ]

TEST_JOINT_ANGLES = dict()
TEST_JOINT_ANGLES['left'] = dict()
TEST_JOINT_ANGLES['right'] = dict()
TEST_JOINT_ANGLES['left']['left_s0'] = 0.0
TEST_JOINT_ANGLES['left']['left_s1'] = 0.0 
TEST_JOINT_ANGLES['left']['left_e0'] = 0.0
TEST_JOINT_ANGLES['left']['left_e1'] = 0.0 
TEST_JOINT_ANGLES['left']['left_w0'] = 0.0
TEST_JOINT_ANGLES['left']['left_w1'] = 0.0
TEST_JOINT_ANGLES['left']['left_w2'] = 0.0
TEST_JOINT_ANGLES['right']['right_s0'] = 0.0 
TEST_JOINT_ANGLES['right']['right_s1'] = 0.0
TEST_JOINT_ANGLES['right']['right_e0'] = 0.0
TEST_JOINT_ANGLES['right']['right_e1'] = math.pi 
TEST_JOINT_ANGLES['right']['right_w0'] = 0.0
TEST_JOINT_ANGLES['right']['right_w1'] = 0.0
TEST_JOINT_ANGLES['right']['right_w2'] = 0.0


def get_joint_angles(user, tfBuffer, test, mirrored):
    """

    @param line: the line described in a list to process
    @param names: joint name keys
    """
    global theta
    joint_angles = dict()
    joint_angles['left'] = dict()
    joint_angles['right'] = dict()

    frames = get_frame_positions(user, tfBuffer)
  
    if frames is None and not test:
        # if there's a problem with tracking, don't move
        return None

    if test:
        # use the hardcoded angles for debugging
        joint_angles = TEST_JOINT_ANGLES
    else:
        # do the math to find joint angles!
        reh = frames['right_hand'] - frames['right_elbow']
        res = frames['right_shoulder'] - frames['right_elbow']
        leh = frames['left_hand'] - frames['left_elbow']
        les = frames['left_shoulder'] - frames['left_elbow']
        rse = np.negative(res)
        lse = np.negative(les)
	
	
        
        # find down vector and normals
        nt = np.cross(frames['right_shoulder'] - frames['torso'], frames['left_shoulder'] - frames['torso'])
        d = np.cross(nt, frames['right_shoulder'] - frames['left_shoulder'])
        rns = np.cross(d, rse) 
        lns = np.cross(d, lse)
        lne = np.cross(leh, les)
        rne = np.cross(reh, res)
        
        # normalize vectors
        reh = reh / np.linalg.norm(reh)
        res = res / np.linalg.norm(res)
        leh = leh / np.linalg.norm(leh)
        les = les / np.linalg.norm(les)
        rse = rse / np.linalg.norm(rse)
        lse = lse / np.linalg.norm(lse)
        nt = nt / np.linalg.norm(nt) 
        d = d / np.linalg.norm(d)
        rns = rns / np.linalg.norm(rns)
        lns = lns / np.linalg.norm(lns)
        lne = lne / np.linalg.norm(lne)
        rne = rne / np.linalg.norm(rne)

        # do the math to find joint angles
        joint_angles['left']['left_s0'] = np.arccos(np.dot(nt,lns)) - math.pi/5.0 # was + 0.0
        joint_angles['left']['left_s1'] = np.arccos(np.dot(d, lse)) - math.pi/2.0
        joint_angles['left']['left_e0'] = np.arccos(np.dot(nt, lne))
        joint_angles['left']['left_e1'] = math.pi - np.arccos(np.dot(leh, les))
        joint_angles['left']['left_w0'] = 0.0
        joint_angles['left']['left_w1'] = 0.0
        joint_angles['left']['left_w2'] = 0.0
        joint_angles['right']['right_s0'] = np.arccos(np.dot(nt,rns)) - math.pi*7.0/8.0 # was - pi
        joint_angles['right']['right_s1'] = np.arccos(np.dot(d, rse)) - math.pi/2.0
        joint_angles['right']['right_e0'] = np.arccos(np.dot(nt, rne)) - math.pi
        joint_angles['right']['right_e1'] = math.pi - np.arccos(np.dot(reh, res))
        joint_angles['right']['right_w0'] = 0.0
        joint_angles['right']['right_w1'] = 0.0
        joint_angles['right']['right_w2'] = theta

    if mirrored:
        return joint_angles
    else:
        # this is the default option, where the teleoperator's
        # left/right arm corresponds to the robot's left/right arm
        unmirrored = dict()
        unmirrored['left'] = dict()
        unmirrored['right'] = dict()
        unmirrored['left']['left_s0'] = joint_angles['right']['right_s0']
        unmirrored['left']['left_s1'] = joint_angles['right']['right_s1']
        unmirrored['left']['left_e0'] = joint_angles['right']['right_e0']
        unmirrored['left']['left_e1'] = joint_angles['right']['right_e1']
        unmirrored['left']['left_w0'] = 0.0 
        unmirrored['left']['left_w1'] = 0.0
        unmirrored['left']['left_w2'] = 0.0
        unmirrored['right']['right_s0'] = joint_angles['left']['left_s0']
        unmirrored['right']['right_s1'] = joint_angles['left']['left_s1']
        unmirrored['right']['right_e0'] = joint_angles['left']['left_e0']
        unmirrored['right']['right_e1'] = joint_angles['left']['left_e1']
        unmirrored['right']['right_w0'] = 0.0
        unmirrored['right']['right_w1'] = 0.0
        unmirrored['right']['right_w2'] = theta
        return unmirrored

def get_frame_positions(user, tfBuffer):
    frame_positions = dict()
    try:
        for frame in FRAMES:
            transformation= tfBuffer.lookup_transform(BASE_FRAME, "%s_%d" % (frame, user), rospy.Time(),rospy.Duration(1.0))
            translation = transformation.transform.translation
            pos = np.array([translation.x, translation.y, translation.z])
            frame_positions[frame] = pos
    except (tf2_ros.LookupException, tf2_ros.ConnectivityException, tf2_ros.ExtrapolationException) as e: 
        print "Problem with kinect tracking."
        print e
        return None
    return frame_positions

def operate_gripper():
    #left = baxter_interface.Gripper('left', CHECK_VERSION)
    right_grip = baxter_interface.Gripper('right', CHECK_VERSION)
    if B_button: 
       print 'button B pressed'
       right_grip.close()
    else:
       right_grip.open()
       print "button B released"

    
def teleoperate(rate, user, test, mirrored):
    """
    Teleoperates the robot based on tf2 frames.

    @param rate: rate at which to sample joint positions in ms

    """
    rate = rospy.Rate(rate)
    # TODO: make these attributes of a class
    tfBuffer = tf2_ros.Buffer()
    listener = tf2_ros.TransformListener(tfBuffer)

    left = baxter_interface.Limb('left')
    right = baxter_interface.Limb('right')


    while not rospy.is_shutdown():
        rate.sleep()
        wii_listen()
#        listener.waitForTransform("/torso_1","camera_depth_frame",now,rospy.Duration(4.0))
        joint_angles = get_joint_angles(user, tfBuffer, test, mirrored)
        
        print joint_angles
        print 'theta = ', theta
        if joint_angles is not None:
            left.set_joint_positions(joint_angles['left'])
            right.set_joint_positions(joint_angles['right'])
            operate_gripper()

    print "Rospy shutdown, exiting loop."
    return True

def main():
    """
    Note: This version of simply drives the joints towards the next position at
    each time stamp. Because it uses Position Control it will not attempt to
    adjust the movement speed to hit set points "on time".
    """
    arg_fmt = argparse.RawDescriptionHelpFormatter
    parser = argparse.ArgumentParser(formatter_class=arg_fmt,
                                     description=main.__doc__)
    parser.add_argument(
        '-r', '--rate', type=int, default=10,
        help='rate to sample the joint positions'
    )

    parser.add_argument(
        '-t', '--test', type=bool, default=False,
        help='use hardcoded test joint angles'
    )
    parser.add_argument(
        '-u', '--user', type=int, default=1,
        help='kinect user number to use'
    )

    parser.add_argument(
        '-m', '--mirrored', type=bool, default=False,
        help='mirror the teleoperators movements'
    )
    args = parser.parse_args(rospy.myargv()[1:])

    print("Initializing node... ")
    rospy.init_node("teleoperation")
    print("Getting robot state... ")
    rs = baxter_interface.RobotEnable(CHECK_VERSION)
    init_state = rs.state().enabled

    def clean_shutdown():
        print("\nExiting...")
        if not init_state:
            print("Disabling robot...")
            rs.disable()
    rospy.on_shutdown(clean_shutdown)

    print("Enabling robot... ")
    rs.enable()

    teleoperate(args.rate, args.user, args.test, args.mirrored)


def wii_listen():
    #rospy.init_node('listener', anonymous=False)
    rospy.Subscriber('/wiimote/state', State , wiimote_state_callback)


def wiimote_state_callback(state_data):
    global theta
    global B_button
    gx = state_data.linear_acceleration_zeroed.x
    gz = state_data.linear_acceleration_zeroed.z
    

    theta = math.atan2(gx,gz)
    
    B_button = state_data.buttons[5]


def declareGlobals():
    global theta
    global B_button
    global right_gripper_state
    theta = math.pi/2.0
    B_button = False
    right_gripper_state = False

if __name__ == '__main__':
    declareGlobals()
    main()