Mapping Robot

# Arvind Kannan, ak633
# Chris Fedors, cjf83
# Final Project
# Mapping Robot
# Last updated: 5/18/17
# mapping_robot.py

import speech_recognition as sr
import random
import datetime
import os
import subprocess
import pyaudio
import time
import picamera
import cv2
import pyaudio
import speech_recognition as sr
import image_detection as imdet
import math
from datetime import datetime

TARGET_AREA = 10000

# GPIO set up
import time
import RPi.GPIO as GPIO

GPIO.setmode(GPIO.BCM)
GPIO.setup(12, GPIO.OUT) # servo
GPIO.setup(16, GPIO.OUT) # servo
GPIO.setup(23, GPIO.IN) # ultrasound
GPIO.setup(24, GPIO.OUT) # ultrasound

GPIO.setup(5, GPIO.IN)
GPIO.setup(6, GPIO.IN)



pulse_start = 0
pulse_end = 0
duration = 0

def get_micro():
	""" Returns current time in microseconds"""
	dt = datetime.now()
	return dt.microsecond

# This callback function is called when GPIO pin 23
# sees a rising or falling edge. 
def ultrasound_receive(channel):
	global pulse_start
	global pulse_end
	global duration
	if(GPIO.input(23) == GPIO.HIGH):
		pulse_start = get_micro()
	else:
		pulse_end = get_micro()
		duration = pulse_end-pulse_start

# Takes ultrasound distance sensor reading, and returns distance
# in cm.
def ultrasound_pulse():
	"""Returns range of objects in centimeters, needs to block for ~100ms"""
	global duration
	print "Pulsing ultrasound"
	# Send a pulse to GPIO pin 24, and look for response
	GPIO.output(24, GPIO.LOW)
	time.sleep(2/1000000.0)
	GPIO.output(24, GPIO.HIGH)
	time.sleep(10/1000000.0)
	GPIO.output(24, GPIO.LOW)
	time.sleep(0.1)
	dist = duration/53.33
	print dist
	if dist <= 0 or dist > 200:
		dist = None
	return dist

# Register callback function ultrasound_receive.
# The bouncetime is small so as to not mess up with timing of the receiving
# of the signal.
GPIO.add_event_detect(23, GPIO.BOTH, callback = ultrasound_receive, bouncetime=1)

# These functions were used in Lab 3 to calculate what 
# commands to send to the motors, and they are still useful here.
def calc_servo_dc(width):
	return (100 * width / (20.0 + width))
def calc_servo_freq(width):
	return (1.0/(.02 + width/1000.0))

# Constants used by the drive function (below).
STAND_STILL_W= 1.5
MAX_CW_W = 1.3
MAX_CCW_W = 1.7

stand_still = calc_servo_dc(STAND_STILL_W)

# Start PWM instances for left and right servos.
p1 = GPIO.PWM(12, calc_servo_freq(STAND_STILL_W))
p2 = GPIO.PWM(16, calc_servo_freq(STAND_STILL_W))
p1.start(stand_still)
p2.start(stand_still)
p = [None, p1, p2]

# This drive function sends the appropriate PWM signal to
# the specified servo.
def drive(num, s):
	d = 0
	f = 50
	if s == 'cw':
		d = calc_servo_dc(MAX_CW_W)
		f = calc_servo_freq(MAX_CW_W)
	elif s == 'ccw':
		d = calc_servo_dc(MAX_CCW_W)
		f = calc_servo_freq(MAX_CCW_W)
	elif s == 's':
		# Send 0 signal to completely stop the servo.
		d = calc_servo_dc(0)
		f = calc_servo_freq(0)
	# num indexes p array, and that allows us to get the 
	# correct PWM instance without needing a conditional here.
	p[num].ChangeDutyCycle(d)
	p[num].ChangeFrequency(f)

# Halt the motors while we're waiting for setup of other
# parts of system.
drive(1, 's')
drive(2, 's')

max_ccw = calc_servo_dc(1.7)
max_cw = calc_servo_dc(1.3)

start = False
target_color = None

# Explore mode is defined as when the robot is looking for 
# any colored target, and will add their positions to the 
# list of targets of that color.
# Mission mode is defined as when the robot is looking for
# a particular colored target. 
EXPLORE_MODE = 'explore'
MISSION_MODE = 'mission'

mode = EXPLORE_MODE

# Intended as a debugging tool (have the robot spin for a while)
# before we started using log files to debug errors.
# Note that the time is not specified. If you call this function,
# you have to time.sleep() afterwards for how long you want the
# robot to spin, and call stop() to stop the motors afterwards.
def spin():
	drive(1, 'cw')
	drive(2, 'cw')


def stop():
	drive(1, 's')
	drive(2, 's')


# Take in the string of what was said to the robot,
# and respond appropriately based on the keywords in the 
# phrase.
def process_command(cmd_str):
	global start
	global target_color
	global mode
	print 'processing command'
	cmd_str = cmd_str.lower()
	toks = cmd_str.split(' ')
	print cmd_str
	if 'spin' in cmd_str:
		spin()
		time.sleep(10)
		stop()
	elif 'start' in cmd_str:
		print '********************STARTING\n\n\n'
		start = True
		mode = EXPLORE_MODE
	elif 'orange' in cmd_str:
		target_color = 'orange'
		mode = MISSION_MODE
	elif 'green' in cmd_str:
		target_color = 'green'
		mode = MISSION_MODE
	elif 'purple' in cmd_str:
		target_color = 'purple'
		mode = MISSION_MODE
	elif 'stop' in cmd_str:
		exit()

# This is called when speech is heard over the background noise.
def callback(recognizer, audio):
	# recognize speech using Google Speech Recognition
	try:
		cmd_str = r.recognize_google(audio)
		print("You said: " + cmd_str)
		process_command(cmd_str)
	# Account for if Google speech recognition is not able to tell what 
	# was said.
	except sr.UnknownValueError:
		print 'Your speech was not recognized'
	except sr.RequestError as e:
		print('You probably lost Internet connection, so you could not get results from Google speech recognition')

# Set up speech recognition instance and the microphone
r = sr.Recognizer()
m = sr.Microphone()
# Account for background noise so the library knows when somebody's speaking near
# the microphone.
with m as source:
	r.adjust_for_ambient_noise(source)

# start listening in the background (note that we don't have to do this inside a `with` statement)
stop_listening = r.listen_in_background(m, callback)
# `stop_listening` is now a function that, when called, stops background listening
print 'finished setting up mic\n\n'

X_PIXELS = 1024
Y_PIXELS = 768


reached_target = False

# Lists with positions of targets of the appropriate color
orange_targets = []
green_targets = []
purple_targets = []


x, y = 0.0, 0.0 # Robot's current x, y position (starts at origin), in cm
robot_orientation = 0.0 # Robot's current orientation (starts facing North)
# For orientation, 0 degrees is facing north, 90 degrees is facing east, etc.

# Given an orientation, return the orientation
# within the range 0 to 359 degrees.
def orientation_in_range(orient):
	while(orient) < 0.0:
		orient += 360.0
	while(orient) >= 360.0:
		orient -= 360.0
	return orient

# Helper function for get_angle function
def calc_ang(O, A):
	return math.degrees(math.atan(float(O) / float(A)))

# If existed in targets list, give the position of that closest one.
# Otherwise, just return the position passed in.
def closest_in_targets_list(target_pos, color):
	TARGET_CLOSENESS_CONSTANT = 70.0
	if color == 'orange':
		targets = orange_targets
	elif color == 'green':
		targets = green_targets
	elif color == 'purple':
		targets = purple_targets
	for elem in targets:
		if abs(target_pos[0] - elem[0]) < TARGET_CLOSENESS_CONSTANT and abs(target_pos[1] - elem[1]) < TARGET_CLOSENESS_CONSTANT:
			return elem
	targets += [target_pos]
	return target_pos

# Get angle between robot and centroid of target.
# This is used so the robot can correct it's course to the target.
def get_angle(rect_center, rect_area):
	global robot_orientation
	#tx, ty = target_pos[0], target_pos[1]
	x_disp = rect_center[0] - 1024/2 # x disp in pixels
	dist_to_object = 21108 / math.sqrt(rect_area)
	ang = calc_ang((26.4 * x_disp / math.sqrt(rect_area)), dist_to_object)

	print 'robot orientation is ' + str(robot_orientation)
	if ang < 0.0:
		return ('ccw', 'ccw', ang)
	return ('cw', 'cw', ang)

# Empirically determined constant based on experimenting with
# rotations of a known angle, like 90 or 180 degrees.
# THis constant times an angle in degrees gives the amount of 
# time in time.sleep() before robot can move on.
ANG_TO_TIME_CONSTANT = 0.005

# Implement correction to robot's course based on horizontal position
# of centroid. Only correct if the angle is substantial, because otherwise
# there would be over-fitting and robot's movement wouldn't be very efficient.
def face_target(rect_center, rect_area):
	global robot_orientation
	direction_servo1, direction_servo2, ang = get_angle(rect_center, rect_area)
	if(abs(ang) > 7.5):
		robot_orientation = ang + robot_orientation
		robot_orientation = orientation_in_range(robot_orientation)
		drive(1, direction_servo1)
		drive(2, direction_servo2)
		print '***Correcting to: ' + str(ang)
		time.sleep(abs(ang) * ANG_TO_TIME_CONSTANT)
		stop()

# Rotate function - update the robot's orientation,
# run the servos in the proper directionm for the correct
# amount of time, then stop the servos.
# We choose the smaller rotation angle, since you could
# of course rotate either way to rotate a certain amount.
def rotate(ang):
	global robot_orientation
	robot_orientation += ang# update
	while robot_orientation > 360.0:
		robot_orientation -= 360.0
	while robot_orientation < 0.0:
		robot_orientation += 360.0
	if ang < 180.0:
		drive(1, 'ccw')
		drive(2, 'ccw')
		time.sleep(ang * ANG_TO_TIME_CONSTANT)
	else:
		drive(1, 'cw')
		drive(2, 'cw')
		ang = 360.0 - ang
		time.sleep(ang * ANG_TO_TIME_CONSTANT)
	stop()

# Turn left and right servos cw and ccw, respectively.
# The code that calls this function is responsible for
# using time.sleep() with the duration of the forward
# movement and calling stop() afterwards.
def drive_forwards():
	drive(1, 'cw')
	drive(2, 'ccw')

# move_towards makes a correction based on the centroid position
# and then drives forwards.
def move_towards(rect_center, rect_area):
	face_target(rect_center, rect_area)
	drive_forwards()

# Return the new position of the robot based on orientation,
# how long the robot moved forward, and the robot's current position.
# As you can tell, the robot's global position (x, y) must be passed 
# in to this function, and the output of this function should be
# assigned to (x, y) to update the global position.
def change_pos(robot_pos, orientation, time_motor):
	new_x = robot_pos[0]
	new_y = robot_pos[1]
	# 21.24 is an empirically defined constant to make
	# the position be in centimeters
	time_motor *= 21.24
	if orientation < 90.0:
		new_x += time_motor * math.sin(math.radians(orientation))
		new_y += time_motor * math.cos(math.radians(orientation))
	elif orientation < 180.0:
		orientation -= 90.0
		new_y -= time_motor * math.sin(math.radians(orientation))
		new_x += time_motor * math.cos(math.radians(orientation))
	elif orientation < 270.0:
		orientation -= 180.0
		new_x -= time_motor * math.sin(math.radians(orientation))
		new_y -= time_motor * math.cos(math.radians(orientation))
	else:
		orientation -= 270.0
		new_x -= time_motor * math.cos(math.radians(orientation))
		new_y += time_motor * math.sin(math.radians(orientation))
	return (new_x, new_y)

# Find the position of the target based on the robot's position and orientation,
# along with the centroid and area of the rectangle corresponding to the target
def get_position(rect_center, rect_area):
	x_disp = rect_center[0] - 1024/2 # x disp in pixels
	# 21108 is a constant determined based on the aspect ratio of the image.
	dist_to_object = 21108 / math.sqrt(rect_area)
	print 'Rect_area is ' + str(rect_area)
	print 'Dist_to_object is ' + str(dist_to_object)

	# 26.4 is a derived constant based on how the trig worked out
	# and on the aspect ratio of the image
	total_dist = math.sqrt(dist_to_object * dist_to_object + (26.4 * x_disp / math.sqrt(rect_area)) * (26.4 * x_disp / math.sqrt(rect_area)))
	dx = total_dist * math.sin(math.radians(robot_orientation + calc_ang((26.4 * x_disp / math.sqrt(rect_area)), dist_to_object)))
	dy = total_dist * math.cos(math.radians(robot_orientation + calc_ang((26.4 * x_disp / math.sqrt(rect_area)), dist_to_object)))
	print 'x distance to target:' + str(dx)
	print 'y distance to target:' + str(dy)
	return (x + dx, y + dy)

# finds closest target of correct color to the robot
def find_closest_target(color):
	if color == 'orange':
		targets = orange_targets
	elif color == 'green':
		targets = green_targets
	elif color == 'purple':
		targets = purple_targets
	min_dist = None
	min_target_pos = None
	for target in targets:
		tx, ty = target[0], target[1]
		dis = math.sqrt((tx - x) * (tx - x) + (ty - y) * (ty - y))
		if min_dist is None:
			min_target_pos = target
			min_dist = dis
		elif dis < min_dist:
			min_target_pos = target
			min_dist = dis
	return min_target_pos

# used for tracking how many consecutive times we've rotated,
# so we can give up and try to move closer
search_count = 0
# used for keeping track of when to override start command being
# given via speech
iter_count = 0
# keep track of the current iteration number we're on
iter2_count = 0
# Start picamera instance
imdet.start_camera()

# The copious numbers of print statements are there because the output of this 
# program was redirected to a file that was used as a log file to help with debugging.
try:
		while True:
			if not start:
				iter_count += 1
				if iter_count == 90000:
					start = True
					print 'TIMER OVERRIDE -- start set to true'
					print 'waiting for you to speak a color'
					a_time = time.time()
					# If 60 seconds have elapsed with a color being recognized
					# by speech recognition library, continue to next iteration.
					while time.time() - a_time < 60.0:
						if target_color is not None:
							break
				continue
			iter2_count += 1
			print '\n\n\nITERATION: ' + str(iter2_count)
			print 'ROBOT_ORIENTATION: ' + str(robot_orientation)
			print 'MODE: ' + mode
			print 'DISTANCE SENSOR READING: ' + str(ultrasound_pulse())
			print 'ORANGE list: '
			print orange_targets
			print 'GREEN list: '
			print green_targets
			print 'PURPLE list: '
			print purple_targets
			print 'ROBOT x and y pos:'
			print str(x) + ', ' + str(y)

			# If the robot is in mission mode, and has not yet reached the target,
			# analyze the image. If the target is not seen, rotate and try to acquire
			# the target. Otherwise, move towards the target, making course corrections
			# as necessary.
			if mode is MISSION_MODE and (not reached_target):
				search_count = 0

				# Find rect of desired color
				rect_center, rect_area, copy_img = imdet.image_detection(target_color)

				# If target of desired color was not found, rotate 15 degrees so we can try again.
				if rect_center is None or rect_area is None:
					result = find_closest_target(target_color)
					rotate(15)
					continue
				# If target of desired color was found, move towards it for a second.
				move_towards(rect_center, rect_area)
				time.sleep(1)
				stop()

				# The distance sensor didn't work too reliably in our range,
				# so we decided not to use it.			
				# DISTANCE_SENSOR_CONSTANT = 20.0 # within 5 cm
				# distance_sensor_val = ultrasound_pulse()

				face_target(rect_center, rect_area)
				# The constants 425000.0 and 157000.0 were empirically determined
				# based on playing around with our image detection test calibration script
				# with our robot at different distances from the target. We observed the
				# area the target took up in the image based on the distance from the target
				# at which we wanted the robot to stop.
				while rect_area < 425000.0:
					# If the target is no longer seen, move back to start mode so we can reset the situation.
					# Otherwise, move towards the target until we're close enough.
					if rect_area < 157000:
						rect_center,rect_area,_ = imdet.image_detection(target_color)
						if rect_center is None or rect_area is None:
							print 'cannot find rect'
							start = False
							target_color = None
							iter_count = 0
						else:
							move_towards(rect_center, rect_area)
					else:
						_, rect_area, _ = imdet.image_detection(target_color, faster = True)
						drive_forwards()
					begin_t = time.time()
					time.sleep(0.5)
					entered_while4 = True
					distance_sensor_val = ultrasound_pulse()
					end_t = time.time()
					x, y = change_pos((x, y), robot_orientation, end_t - begin_t)
					print 'INRANGE'
				stop()
				print '===REACHED DESIRED TARGET'
				rotate(180)
				start = False
				target_color = None
				iter_count = 0

			elif mode is MISSION_MODE:
				# The target has been reached, so we switch back to explore mode (and thus require 
				# speech instruction to switch back to mission mode)
				mode = EXPLORE_MODE
				search_count = 0
			# If the mode is EXPLORE_MODE, find targets of any of the three colors
			elif mode is EXPLORE_MODE:
				# Find targets by looking for large rectangles of the correct color
				# This is accomplished by calling the image_detection module's
				# namesake function, image_detection.
				rect_centerO, rect_areaO, copy_imgO = imdet.image_detection('orange')
				rect_centerG, rect_areaG, copy_imgG = imdet.image_detection('green')
				rect_centerP, rect_areaP, copy_imgP = imdet.image_detection('purple')
				print 'IMAGE DETECTION DONE'
				print rect_areaG
				if rect_areaO is not None:
					rel_distO = float(TARGET_AREA) / float(rect_areaO)
					print 'rel distance to orange: ' + str(rel_distO)
				if rect_areaG is not None:
					rel_distG = float(TARGET_AREA) / float(rect_areaG)
					print 'rel distance to green: ' + str(rel_distG)
				if rect_areaP is not None:
					rel_distP = float(TARGET_AREA) / float(rect_areaP)
					print 'rel distance to purple: ' + str(rel_distP)

				# A constant -- the minimum pixel area of the polygon in the image
				# for which we consider it to be the target.
				MIN_ACCEPTABLE_AREA = 5000

				if rect_areaO is not None and rect_areaO > MIN_ACCEPTABLE_AREA:
					target_posO = get_position(rect_centerO, rect_areaO)
					closest_in_targets_list(target_posO, 'orange')
					print '**** Orange rect found\n\n\n'
					if rect_centerO[0] < 1024/2:
						print 'Orange rect is left of center from robot perspective by: ' + str(1024/2 - rect_centerO[0]) + ' pixels'
					else:
						print 'Orange rect is right of center from robot perspective by: ' + str(rect_centerO[0] - 1024/2) + ' pixels'
					print 'orange target position is: ' + str(target_posO)
					move_towards(rect_centerO, rect_areaO)
					time.sleep(1)
					x, y = change_pos((x, y), robot_orientation, 1.0)
					stop()
					# The distance sensor didn't work too reliably in our range,
					# so we decided not to use it.			
					# DISTANCE_SENSOR_CONSTANT = 20.0 # within 5 cm
					# distance_sensor_val = ultrasound_pulse()
					entered_while1 = False
					face_target(rect_centerO, rect_areaO)
					# The constants 425000.0 and 157000.0 were empirically determined
					# based on playing around with our image detection test calibration script
					# with our robot at different distances from the target. We observed the
					# area the target took up in the image based on the distance from the target
					# at which we wanted the robot to stop.
					while rect_areaO < 425000.0:
						# If the target is no longer seen, move back to start mode so we can reset the situation.
						# Otherwise, move towards the target until we're close enough.
						print 'area: ' + str(rect_areaO)
						if rect_areaO < 157000.0:
							rect_centerO,rect_areaO,_ = imdet.image_detection("orange")
							if rect_centerO is None or rect_areaO is None:
								print 'cannot find rect'
								start = False
								target_color = None
								iter_count = 0
							else:
								move_towards(rect_centerO, rect_areaO)
						else:
							_, rect_areaO, _ = imdet.image_detection("orange", faster = True)
							drive_forwards()
							time.sleep(0.5)
						begin_t = time.time()
						distance_sensor_val = ultrasound_pulse()
						end_t = time.time()
						x, y = change_pos((x, y), robot_orientation, end_t - begin_t)
					stop()

					print '===REACHED ORANGE TARGET'
					rotate(180)
					start = False
					target_color = None
					iter_count = 0

				elif rect_areaG is not None and rect_areaG > MIN_ACCEPTABLE_AREA:
					target_posG = get_position(rect_centerG, rect_areaG)
					closest_in_targets_list(target_posG, 'green')
					print '**** Green rect found\n\n\n\n'
					print 'green target position is: ' + str(target_posG)
					# Move towards the green target that was found
					move_towards(rect_centerG, rect_areaG)
					time.sleep(1)
					x, y = change_pos((x, y), robot_orientation, 1.0)
					stop()
					# The distance sensor didn't work too reliably in our range,
					# so we decided not to use it.			
					# DISTANCE_SENSOR_CONSTANT = 20.0 # within 5 cm
					# distance_sensor_val = ultrasound_pulse()
					face_target(rect_centerG, rect_areaG)
					# The constants 425000.0 and 157000.0 were empirically determined
					# based on playing around with our image detection test calibration script
					# with our robot at different distances from the target. We observed the
					# area the target took up in the image based on the distance from the target
					# at which we wanted the robot to stop.
					while rect_areaG < 425000.0:
						# If the target is no longer seen, move back to start mode so we can reset the situation.
						# Otherwise, move towards the target until we're close enough.
						if rect_areaG < 157000.0:
							rect_centerG,rect_areaG,_ = imdet.image_detection("green")
							if rect_centerG is None or rect_areaG is None:
								print 'cannot find rect'
								start = False
								target_color = None
								iter_count = 0
							else:
								move_towards(rect_centerG, rect_areaG)
						else:
							_, rect_areaG, _ = imdet.image_detection("green", faster = True)
							drive_forwards()
						begin_t = time.time()
						time.sleep(0.5)
						distance_sensor_val = ultrasound_pulse()
						end_t = time.time()
						x, y = change_pos((x, y), robot_orientation, end_t - begin_t)
					stop()

					print '===REACHED GREEN TARGET'
					rotate(180)
					start = False
					target_color = None
					iter_count = 0
				elif rect_areaP is not None and rect_areaP > MIN_ACCEPTABLE_AREA:
					target_posP = get_position(rect_centerP, rect_areaP)
					closest_in_targets_list(target_posP, 'purple')
					print '**** Purple rect found\n\n\n\n'
					print 'purple target position is: ' + str(target_posP)
					move_towards(rect_centerP, rect_areaP)
					time.sleep(1)
					x, y = change_pos((x, y), robot_orientation, 1.0)

					stop()
					# The distance sensor didn't work too reliably in our range,
					# so we decided not to use it.			
					# DISTANCE_SENSOR_CONSTANT = 20.0 # within 5 cm
					# distance_sensor_val = ultrasound_pulse()
					face_target(rect_centerO, rect_areaO)
					# The constants 425000.0 and 157000.0 were empirically determined
					# based on playing around with our image detection test calibration script
					# with our robot at different distances from the target. We observed the
					# area the target took up in the image based on the distance from the target
					# at which we wanted the robot to stop.
					while rect_areaP < 425000.0:
						# If the target is no longer seen, move back to start mode so we can reset the situation.
						# Otherwise, move towards the target until we're close enough.
						if rect_areaP < 157000:
							rect_centerP,rect_areaP,_ = imdet.image_detection("purple")
							if rect_centerP is None or rect_areaP is None:
								print 'cannot find rect'
								start = False
								target_color = None
								iter_count = 0
							else:
								move_towards(rect_centerP, rect_areaP)
						else:
							_, rect_areaP, _ = imdet.image_detection("purple", faster = True)
							drive_forwards()
						begin_t = time.time()
						time.sleep(0.5)
						distance_sensor_val = ultrasound_pulse()
						end_t = time.time()
						x, y = change_pos((x, y), robot_orientation, end_t - begin_t)
					stop()

					print '===REACHED PURPLE TARGET'
					rotate(180)
					start = False
					target_color = None
					iter_count = 0
				else:
					# Try rotating to look for a better direction.
					# If that doesn't work for 360 degrees of rotation,
					# move forward for a second, and try again.
					print '*** could not find anything -- ROTATING'
					ROTATION_CONSTANT = 30
					rotate(ROTATION_CONSTANT) # rotate
					search_count += 1
					if search_count == 360 / ROTATION_CONSTANT:
						search_count = 0
						drive_forwards()
						time.sleep(1)
						x, y = change_pos((x, y), robot_orientation, 1.0)
						stop()
					continue
# When testing with a keyboard attached, the script is terminated with
# Ctrl-C, since it didn't make sense to have a script that lasted a particular
# amount of time.
except KeyboardInterrupt:
	pass
# Cleanup - closing picamera instance, stopping PWM instances, GPIO cleanup
imdet.end_camera()
p1.stop()
p2.stop()
GPIO.cleanup()

# Arvind Kannan, ak633
# Chris Fedors, cjf83
# Final Project - mapping robot
# Last modified: 5/18/17
# image_detection.py (target recognition module)

import picamera
from picamera.array import PiRGBArray
import numpy
import numpy as np
import cv2
import time

camera = 0
up_close = False

#Thresolds for colors in HSV
lower_grey = np.array([0,0,0])
upper_grey = np.array([255,255,255])

lower_purple = np.array([120,20,80])
upper_purple = np.array([170,255,220])

lower_orange = np.array([0, 70, 70])
upper_orange = np.array([50, 255, 255])

lower_green = np.array([50,60,40])
upper_green = np.array([95,255,255])

def start_camera():
	global camera
	camera = picamera.PiCamera()
	camera.resolution = (1024, 768)

def end_camera():
	global camera
	camera.close()


def get_color_mask(color, hsv):
	if color == "purple":
		mask = cv2.inRange(hsv, lower_purple, upper_purple)
	elif color == "orange":
		mask = cv2.inRange(hsv, lower_orange, upper_orange)
	elif color == "green":
		mask = cv2.inRange(hsv, lower_green, upper_green)
	else:
		mask = cv2.inRange(hsv, lower_grey, upper_grey)
		print "unrecognized color"
	return mask

def take_picture(camera):
	rawCapture = PiRGBArray(camera)
	time.sleep(2)
	# grab an image from the camera
	camera.capture(rawCapture, format="bgr")
	img = rawCapture.array
	return img

def convert_BW(hsv, img, color):
	mask = get_color_mask(color, hsv)
	per = cv2.countNonZero(mask)
	per = per/(1024.0*768.0)
	print str(color) + " is %" + str(per) + " of image"
	res= cv2.bitwise_and(img,img,mask=mask)
	gimg = cv2.cvtColor(res, cv2.COLOR_BGR2GRAY)
	return gimg, per

def process_img(gimg):
	gimg = cv2.GaussianBlur(gimg,(19,19),0)
	gimg = cv2.threshold(gimg, 25, 255, cv2.THRESH_BINARY)[1]
	return gimg

def find_edges(gimg):
	eimg = cv2.Canny(gimg, 50, 150)
	# Following two lines were commented out because they caused
	# a Gtk display error when not run in the RPi desktop, but when
	# in the desktop, this imshow was very helpful for debugging.
	#cv2.imshow('edged', eimg)
	#cv2.waitKey(0)
	return eimg

def find_contours(eimg,img):
	pimg = img.copy()
	squares = []
	maxsquare = None
	maxsquarearea = 0
	maxcontour = None
	contours, hierarchy = cv2.findContours(eimg, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
	for contour in contours:
		# Approximate the corners of the polygon
		area = cv2.contourArea(contour)
		perimeter = cv2.arcLength(contour, True)
		corners = cv2.approxPolyDP(contour, perimeter * 0.05, True)

		#if len(corners) != 4:
		if len(corners) > 5 or len(corners) < 2:
			continue
		# Ignore concave shapes
		if not cv2.isContourConvex(corners):
			continue
		if area < 1000:
			continue

		# Find the quad of largest area
		corners = corners.reshape((corners.shape[0], 2))
		squares.append(corners)
		if area > maxsquarearea:
			maxsquare = corners
			maxsquarearea = area
			maxcontour = contour
			#Calculate centroid
			M = cv2.moments(contour)
			x_center = int(M['m10']/M['m00'])
			y_center = int(M['m01']/M['m00'])

	if maxsquare is None:
		print 'no square with area > 200'
		x_center = None
		y_center = None
		maxsquarearea = None
	else:
		# Draw outlines on the image
		cv2.circle(pimg, (x_center,y_center),10,(255, 0, 0),3)
		cv2.polylines(pimg, squares, True, (0, 255, 255), 5)
	return (x_center,y_center), maxsquarearea, pimg

def image_detection(color, faster = False):
	"""Takes in image color (purple, orange, or green)
	Returns (x,y), area, image
	-(x,y) are x and y coordinates of centriod of largest recatangle in image
	-area is area of largest rectangle
	-image has largest reactangle highlighted and centriod circled
	"""
	global camera
	print '***Camera set up'
	img = take_picture(camera)
	hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
	gimg, per = convert_BW(hsv,img,color)
	if(per < 0.20 and not faster):
		gimg = process_img(gimg)
		print '***Filtered Image'
		eimg = find_edges(gimg)
		print '***Contoured Image'
		return find_contours(eimg,img)
	else:
		print "UP CLOSE"
		return (1024/2.0,768/2.0), per*(1024*768),img

# Arvind Kannan, ak633
# Chris Fedors, cjf83
# Final Project - mapping robot
# Last modified: 5/18/17
# img_detection_test.py (calibration script for checking target recognition)

import cv2
from image_detection import *

start_camera()

data = image_detection("orange")
print data[1]
cv2.imshow("Orange", data[2])
cv2.waitKey(0)
data = image_detection("purple")
print data[1]
cv2.imshow("Purple", data[2])
cv2.waitKey(0)
data = image_detection("green")
print data[1]
cv2.imshow("Green", data[2])
cv2.waitKey(0)


end_camera()