Robotic Arm controlled by Raspberry Pi

robotarm.py

import RPi.GPIO as GPIO
import time
import pygame
import sys

##############################################
#
# ROBOTIC ARM CONTROLLER
#
# The robotic arm has five motors
# Each motor is being controlled by an L293D or an L298N driver
# 2 RPi GPIO pins are required per motor - so 10 in total
#
# The motors are powered by 4 D cells in the original robotic arm box
# The RPi and the motor drivers are all connected to a common GND
#
##############################################

# a motor object
# one of these is needed for each of the robotic arm motors
class Motor:
    def __init__(self, motorId, pinRight, pinLeft, keyRight, keyLeft):
        self.pinRight = pinRight
        self.pinLeft = pinLeft
        self.keyRight = keyRight
        self.keyLeft = keyLeft
        self.motorId = motorId
        self.setup()
    def setup (self):
        GPIO.setup(self.pinRight,GPIO.OUT)
        GPIO.setup(self.pinLeft,GPIO.OUT)
    def control (self,key):
        if key == self.keyRight:
            GPIO.output(self.pinRight,True)
            GPIO.output(self.pinLeft,False)
        elif key == self.keyLeft:
            GPIO.output(self.pinRight,False)
            GPIO.output(self.pinLeft,True)
        else:
            GPIO.output(self.pinRight,False)
            GPIO.output(self.pinLeft,False)

# key handler
# keeps a list indexed by key
# each key uniquely identifies the motor responsible for it
# the list is setup after the motor objects are all defined
class KeyHandler:
    def __init__(self):
        self.handler = {}
    def getMotor(self,key):
        if key in self.handler:
            return self.handler[key]
        else:
            return -1
    def setup (self):
        for i in range (0,len(motors)):
            self.handler[motors[i].keyRight] = i
            self.handler[motors[i].keyLeft] = i

# make sure all the motors are in a known, stopped state
def stopAllMotors ():
    for motor in range (0,MAX_MOTORS):
        motors[motor].control(stop)    

######################################
#        
# main part of the code
#
######################################

# when a key is released, a stop event is
# sent to the relevant motor rather than the key number
stop = 0

# unique identifiers for each motor
motorJaws = 0
motorShoulder = 1
motorElbow = 2
motorWrist = 3
motorBase = 4
MAX_MOTORS = 5

GPIO.setmode(GPIO.BCM)

# create all 5 motor objects
motors = []
motors.append(Motor(motorJaws,23,24,pygame.K_RIGHT,pygame.K_LEFT))
motors.append(Motor(motorShoulder,7,8,pygame.K_UP,pygame.K_DOWN))
motors.append(Motor(motorElbow,18,25,pygame.K_a,pygame.K_d))
motors.append(Motor(motorWrist,14,15,pygame.K_w,pygame.K_x))
motors.append(Motor(motorBase,9,10,pygame.K_j,pygame.K_k))

# create the key handler
keyHandler = KeyHandler()
# and populate it with information from the motor objects
keyHandler.setup()

# create a pygame window so that we can accept keyboard events
# without having to use python's basic raw_input function
pygame.init()
screen = pygame.display.set_mode([640,480])
background = pygame.Surface(screen.get_size())
background.fill([255,255,255])

# check for keyboard events until the user decides to quit
while True:

    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            stopAllMotors()
            sys.exit()
     
        # User pressed down on a key
        if event.type == pygame.KEYDOWN:
            motor = keyHandler.getMotor(event.key)
            if motor != -1:
                motors[motor].control(event.key)
                
        # User let up on a key
        if event.type == pygame.KEYUP:
            motor = keyHandler.getMotor(event.key)
            if motor != -1:
                motors[motor].control(stop)