# Finds a good threshold for finding reflective tape
# illuminated with green LEDs.
# The position of the tape is then displayed using
# the center-of-mass method.
# This method is simple, but might be fooled by imperfect
# calibration or another robot with green lights.
# http://www.pyimagesearch.com/2014/08/04/opencv-python-color-detection/

import pylab
import numpy
from picamera.array import PiRGBArray
from picamera import PiCamera
import cv2

camera = PiCamera()
camera.resolution = (1280/4, 768/4)
rawCapture = PiRGBArray(camera)
camera.capture(rawCapture, format = 'bgr')
image = rawCapture.array
cv2.imwrite('imagebase.png',image)

# Try different thresholds for green,
# and look at the pictures to see if you like.
for thresh in [0,64,128,192,230,240,250]:
	lower = numpy.array([0,thresh,0],dtype = "uint8")
	upper = numpy.array([255,255,255],dtype = "uint8")
	# Set green below thresh to zero
	mask = cv2.inRange(image, lower, upper)
	output = cv2.bitwise_and(image,image,mask=mask)
	#cv2.imshow(str(thresh), numpy.hstack([image,output]))
	cv2.imshow('Threshold = ' + str(thresh), output)
	cv2.imwrite('imagethresh'+str(thresh)+'.png',output)
	cv2.waitKey(0)

# Print the bright pixels.
# Ignore pixels that are far away from the main blob.
print 'Bright pixels that are hopefully your reflective tape!'
print 'row index, column index, color triplet'
for jrow in range(768/4):
	for jcol in range(1280/4):
		if (mask[jrow,jcol]):
			print jrow,jcol,image[jrow,jcol,:]

# These values were found using the above loop.
# For experts, they could be found statistically.
# Notice the green is so bright it is saturating
# red and blue, so it appears almost white.
lower = numpy.array([241,251,196],dtype = "uint8")
upper = numpy.array([255,255,255],dtype = "uint8")

# Find the location of the tape blob for the image we
# have been using for thresholding.
# Now do it dynamically, so you can move the camera
# and/or the tape and see it's measured position change.
while True:
	mask = cv2.inRange(image, lower, upper)
	output = cv2.bitwise_and(image,image,mask=mask)
	nz = numpy.nonzero(output[:,:,1])
	fracx = numpy.mean(nz[1])/(1280/4)
	fracy = numpy.mean(nz[0])/(768/4)
	print fracx, fracy
	cv2.imshow(str(round(fracx,2))+' '+str(round(fracy,2)), output)
	cv2.imwrite('loc_'+str(round(fracx,2))+'_'+str(round(fracy,2))+'.png',output)
	cv2.waitKey(0)
	# Take the next pic
	camera.resolution = (1280/4, 768/4)
	rawCapture = PiRGBArray(camera)
	camera.capture(rawCapture, format = 'bgr')
	image = rawCapture.array