kbd_flight.py

#!/usr/bin/python2
import curses
import sys
import time

import pygame
try:
    import curses
except:
    curses = False
# from evdev import InputDevice, list_devices, ecodes

import drone


DASHBOARD = """
 Flight control for Eachine/JJRC

            throttle                            pitch

               [w]                               [up]

              {throttle_u:>3}%                               {pitch_u:>3}%

 yaw [a] {yaw_l:>3}%      {yaw_r:>3}%  [d]         [left] {roll_l:>3}%      {roll_r:>3}%  [right]  roll

              {throttle_d:>3}%                               {pitch_d:>3}%

               [s]                              [down]


roll:{roll:>3} pitch:{pitch:>3} throttle:{throttle:>3} yaw:{yaw:>3}
pressed: {pressed_keys}
Command for drone: 0x{hex_command}

Unix time: {unixtime}

Press Ctrl+C to exit
"""

SUPPORTED_JOYSTICKS = {
  'DragonRise Inc.   Generic   USB  Joystick  ': {
    'buttons': {
      'BTN_START': 9,
      'BTN_SELECT': 8,
      'BTN_L1': 6,
      'BTN_L2': 4,
      'BTN_R1': 7,
      'BTN_R2': 5,
      'BTN_1': 0,
      'BTN_2': 1,
      'BTN_3': 2,
      'BTN_4': 3,
    },
    'axis': {
      'LEFT_VERTICAL': (3, -1),
      'LEFT_HORIZONTAL': (0, 1),
      'RIGHT_VERTICAL': (4, -1),
      'RIGHT_HORIZONTAL': (1, 1),
    }
  },
  
  'xiaoji Gamesir-G4s 2.0b': {
    'buttons': {
      'BTN_START': 11,
      'BTN_SELECT': 10,
      'BTN_L1': 6,
      'BTN_L2': 8,
      'BTN_R1': 7,
      'BTN_R2': 9,
      'BTN_1': 0,
      'BTN_2': 1,
      'BTN_3': 3,
      'BTN_4': 4,
    },
    'axis': {
      'LEFT_VERTICAL': (1, -1),
      'LEFT_HORIZONTAL': (0, 1),
      'RIGHT_VERTICAL': (3, -1),
      'RIGHT_HORIZONTAL': (2, 1),
    }
  },
}


def detect_joystick():
    sys.stdout.write('{} joysticks found.\n'.format(str(pygame.joystick.get_count())))
    if pygame.joystick.get_count() > 0:
        j = pygame.joystick.Joystick(0)
        j.init()
        sys.stdout.write('found: ' + j.get_name())
        return j
    return None


def wait_for_joystick():
    joystick_count = pygame.joystick.get_count()
    if joystick_count > 0:
        for i in range(joystick_count):
            j = pygame.joystick.Joystick(i)
            j.init()
            sys.stdout.write('found joystick "{}"'.format(j.get_name()))
            if j.get_name() in SUPPORTED_JOYSTICKS.keys():
                return j
    return None


def init_screen():
    screen = curses.initscr()
    curses.noecho()
    # set getch() non-blocking
    screen.nodelay(1)
    # escape sequences generated by some keys (keypad, function keys) will be interpreted by curses
    screen.keypad(1)
    screen.clearok(1)
    return screen


def redraw_screen(screen, roll, pitch, throttle, yaw,
                  pressed=None, drone_command=''):
    if pressed is None:
        pressed = []

    state = {
        'throttle_u': int((throttle+1)/2.0*100),
        'throttle_d': int((1-(throttle+1)/2.0)*100),
        'yaw_r': int((yaw+1)/2.0*100),
        'yaw_l': int((1-(yaw+1)/2.0)*100),
        'pitch_u': int((pitch+1)/2.0*100),
        'pitch_d': int((1-(pitch+1)/2.0)*100),
        'roll_r': int((roll+1)/2.0*100),
        'roll_l': int((1-(roll+1)/2.0)*100),
        'unixtime': time.time(),
        'hex_command': drone_command.encode('hex'),
        'pressed_keys': str(pressed),
        'roll': roll,
        'pitch': pitch,
        'throttle': throttle,
        'yaw': yaw,
    }

    screen.clear()

    for i, line in enumerate(DASHBOARD.split('\n')):
        screen.addstr(i, 0, line.format(**state))

    screen.refresh()


def parse_joystick_input(joystick):
    throttle = 0
    yaw = 0
    pitch = 0
    roll = 0
    commands = set()
    pressed = set()

    j_settings = SUPPORTED_JOYSTICKS[joystick.get_name()]

    j_axis = j_settings['axis']
    throttle = joystick.get_axis(j_axis['LEFT_VERTICAL'][0])*j_axis['LEFT_VERTICAL'][1]
    yaw = joystick.get_axis(j_axis['LEFT_HORIZONTAL'][0])*j_axis['LEFT_HORIZONTAL'][1]
    pitch = joystick.get_axis(j_axis['RIGHT_VERTICAL'][0])*j_axis['RIGHT_VERTICAL'][1]
    roll = joystick.get_axis(j_axis['RIGHT_HORIZONTAL'][0])*j_axis['RIGHT_HORIZONTAL'][1]

    j_buttons = j_settings['buttons']
    for btn, code in j_buttons.items():
        if joystick.get_button(code) == 1:
            pressed.add(btn)

    if 'BTN_START' in pressed:
        commands.add('spin_up')
    if 'BTN_SELECT' in pressed:
        commands.add('shut_off')
    if 'BTN_R1' in pressed or 'BTN_R2' in pressed:
        commands.add('force')

    return roll, pitch, throttle, yaw, commands, pressed


DEBUG = True
if DEBUG:
    MAX_POWER = 0.5


def main_loop(drone1, screen=None, kbd=None, joystick=None):
    roll, pitch, throttle, yaw = 0, 0, 0, 0
    cmd = None

    joystick = wait_for_joystick()
    if joystick is None:
        time.sleep(1)
        sys.stdout.write('No supported joysticks found, exiting\n')
        sys.exit(1)
    try:
        while 1:
            roll, pitch, throttle, yaw, commands, pressed = parse_joystick_input(joystick)

            max_power = 1.0 if 'force' in commands else MAX_POWER

            if 'shut_off' in commands:
                cmd = 'shut_off'
            elif 'land' in commands:
                cmd = 'land'
            elif 'spin_up' in commands:
                cmd = 'spin_up'
            elif 'calibrate' in commands:
                cmd = 'calibrate'
            else:
                cmd = None

            yaw *= max_power
            pitch *= max_power
            roll *= max_power

            drone_command = drone.get_command_string(
                roll=roll, pitch=pitch, throttle=throttle, yaw=yaw, command=cmd,
                altitude_hold=True
            )
            drone1.execute(drone_command)

            pygame.event.pump()

            if screen is not None:
                redraw_screen(screen,
                              roll, pitch, throttle, yaw,
                              drone_command=drone_command,
                              pressed=pressed)
            clock.tick(20)
    finally:
        drone1.disconnect()
        curses.endwin()


if __name__ == "__main__":
    drone1 = drone.Drone()
    pygame.init()
    pygame.joystick.init()
    # if you are running script without TTY, don't install curses in virtualenv
    screen = init_screen() if curses else None
    kbd = None
    clock = pygame.time.Clock()
    main_loop(drone1, screen)