Python script to control the Donatello Robot. Uses WiringPi and OpenCV.

Donatello.py

# Marc-Daniel Julien
# Akiva Krauthamer
# Intro to Physical Compting
# Carnegie Mellon University

import cv2
import time, sys
import numpy as np
from scipy import misc
import wiringpi2 as wpi
import threading
import signal

# Globals
THRESH = 100
LOWER_BLUE = np.array([180*(135.0/360.0),   100, 100], dtype=np.uint8)
HIGHER_BLUE = np.array([180*(270.0/360.0),   255, 255], dtype=np.uint8)
DISPLAY = False

INPUT = 0
OUTPUT = 1
PWM = 2

HIGH = 1
LOW = 0

# RPi Pins
SPEED_PIN = 18
DIR_PIN = 24
INA_PIN = 4
INB_PIN = 17

# Wheel states
ROLLING = 0
STRAIGHT = 1
STOP = 2

# Wheel direction
LEFT = 0
RIGHT = 1
FORWARD = 2
REVERSE = 3

STRAIGHT_STEPS = [LEFT, FORWARD, RIGHT, FORWARD]*10
ROLLING_STEPS = [LEFT, LEFT, LEFT, LEFT]*10

# Interval between each step
INTERVAL = 0.1

# Webcam object
capture = None

# Control Object
ct = None

# A Class used to control the motion of the robot
class ControlThread():
	def __init__(self):
		self.state = ROLLING
		self.steps = ROLLING_STEPS
		self.stepn = 0
		self.counter = 0
		self.done = False

	def run(self):
		while(not self.done):
			time.sleep(INTERVAL)
			if(self.state == STRAIGHT):
				print "Straight"
			if(self.state == ROLLING):
				print "Rolling"
			if(self.state == STOP):
				print "Stop"
			
			self.next_step()

    # While iterating through the steps in the list
    # This function will go in the direction indicated
	def next_step(self):
		if(self.stepn < len(self.steps)):
			if(self.steps[self.stepn] == LEFT):
				self.turn_left()
			elif(self.steps[self.stepn] == RIGHT):
				self.turn_right()
			elif(self.steps[self.stepn] == FORWARD):
				self.forward()
			elif(self.steps[self.stepn] == REVERSE):
				self.reverse()
				
    # Called when blue is detected,
    # switches the robot to Straight Mode
	def blue_signal(self):
		print "Blue!"
		if(self.state == ROLLING):
			self.state = STRAIGHT
			self.steps = STRAIGHT_STEPS
			self.stepn = 0

    # Spin wheel 2 right
	def w2right(self):
		wpi.digitalWrite(INA_PIN, HIGH)
		wpi.digitalWrite(INB_PIN, LOW)
    # Spin wheel 2 left
	def w2left(self):
		wpi.digitalWrite(INA_PIN, LOW)
		wpi.digitalWrite(INB_PIN, HIGH)
    # Spin wheel 1 right slow
	def w1speed1r(self):
		wpi.softPwmWrite(SPEED_PIN, 50)
		wpi.digitalWrite(DIR_PIN, LOW)
	# Spin wheel 1 right fast
	def w1speed2r(self):
		wpi.softPwmWrite(SPEED_PIN, 100)
		wpi.digitalWrite(DIR_PIN, LOW)
    # Spin wheel 1 left slow
	def w1speed1l(self):
		wpi.softPwmWrite(SPEED_PIN, 50)
		wpi.digitalWrite(DIR_PIN, HIGH)
	# Spin wheel 1 left fast	
	def w1speed2l(self):
		wpi.softPwmWrite(SPEED_PIN, 0)
		wpi.digitalWrite(DIR_PIN, HIGH)

    # Causes the robot to turn left
	def turn_left(self):
		self.w2left()
		self.w1speed1r()
    # Causes the robot to turn right
	def turn_right(self):
		self.w2right()
		self.w1speed1l()
    # Causes the robot to go forward
	def forward(self):
		self.w2right()
		self.w1speed2r()
    # Causes the robot to go backward
	def reverse(self):
		self.self.w2left()
		self.w1speed2l()

	def stop(self):
		wpi.digitalWrite(INA_PIN, HIGH)
		wpi.digitalWrite(INB_PIN, HIGH)
		wpi.softPwmWrite(SPEED_PIN, 100)
		wpi.digitalWrite(DIR_PIN, HIGH)


	def end(self):
		self.stop()
		self.done = True

# Counts the percentage of the image that contains a marker ink
def percent_marker(imgarr):
	count = 0
	for val in imgarr.flatten():
		if(val > THRESH):
			count+=1
	return float(count)/(imgarr.shape[0]*imgarr.shape[1])

# Captures and returns a binary image of the marker ink
def get_frame(capture):
	(ret, frame) = capture.read()
	frame = cv2.blur(frame, (5,5))
	frame = cv2.resize(frame, (32,32), interpolation = cv2.INTER_CUBIC)
	frame = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
	frame = cv2.inRange(frame, LOWER_BLUE, HIGHER_BLUE)
	return frame



def main():
	global capture
	global ct
	global INTERVAL
	global DISPLAY
	
	# Interval for each step in wheel motion
	INTERVAL = float(sys.argv[1])	
	
	# Displays the webcam feed, for debugging
	if(len(sys.argv) > 2):
		DISPLAY  = True

    # Set up WiringPi Pins
    # One set of pins controlls a wheel at various speeds
    # One set of pins controls a motor with only digitial i/o
	wpi.wiringPiSetupGpio()
	wpi.pinMode(DIR_PIN, OUTPUT)
	wpi.softPwmCreate(SPEED_PIN, 0, 100)
	wpi.pinMode(INA_PIN, 1)
	wpi.pinMode(INB_PIN, 1)

	# Get Webcam feed
	capture = cv2.VideoCapture(-1)
	if(not capture.isOpened()):
		print "No Cam found, exiting"
		exit()

	# Start controller class
	ct = ControlThread()
	ct.start_wheel()

	# Begin main loop
	print "Starting"
	while True:
	    # Execute next step for wheel
		ct.run()
		
		# Grab frame only when in rolling state
		if(ct.state == ROLLING):
    		try:
    			frame = get_frame(capture)
    			percent = percent_marker(frame)
    		except Exception, e:
    			print e
    			continue;

        # If there's any blue detected
        # Signal wheel to change state
		if (percent > 0.0):
			ct.blue_signal()

        # Display original webcam feed
        # and the processed image for detecting blue color
		if(DISPLAY):
			cv2.imshow("frame", frame);
			cv2.imshow("frame2", capture.read()[1])
		
		if(cv2.waitKey(20)&0xFF == ord('q')):
			break

    # Close camera on exit
	capture.release()


# Releases the Webcam resources and stops the wheel
# before quitting
def handler(signum, frame):
	print "Killed"
	capture.release()
	ct.end()	

if __name__ == "__main__":
    # Install Signal Handlers for quitting
	signal.signal(signal.SIGINT, handler)
	signal.signal(signal.SIGTERM, handler)
	
	# Run main program
	main()