import Tkinter as tk
import cv2, sys, time, os, math
from PIL import Image, ImageTk
import numpy as numpy

from os import listdir
from os.path import isfile, join

import RPi.GPIO as GPIO

# Load the BCM V4l2 driver for /dev/video0
os.system('sudo modprobe bcm2835-v4l2')
# Set the framerate ( not sure this does anything! )
os.system('v4l2-ctl -p 4')

width, height = 320, 240
cap = cv2.VideoCapture(0)
cap.set(cv2.cv.CV_CAP_PROP_FRAME_WIDTH,  width)
cap.set(cv2.cv.CV_CAP_PROP_FRAME_HEIGHT, height)

cascPath = '/home/pi/lbpcascade_frontalface.xml'
faceCascade = cv2.CascadeClassifier(cascPath)

root = tk.Tk()
root.attributes("-fullscreen", True)
root.bind('<Escape>', lambda e: root.quit())

lmain = tk.Label(root)
lmain.pack()

last_image_faces = []
saved = False

GPIO.setmode(GPIO.BOARD)
GPIO.setup(12, GPIO.IN, pull_up_down = GPIO.PUD_UP)
GPIO.setup(16, GPIO.IN, pull_up_down = GPIO.PUD_UP)
GPIO.setup(18, GPIO.IN, pull_up_down = GPIO.PUD_UP)



def show_frame():
    _, frame = cap.read()
    frame = cv2.flip(frame, 1)
    frame = faceDetect(frame)
    #frame = cv2.resize(frame, (320,240))
    buttonPress()
    cv2image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGBA)
    img = Image.fromarray(cv2image)
    imgtk = ImageTk.PhotoImage(image=img)
    lmain.imgtk = imgtk
    lmain.configure(image=imgtk)
    lmain.after(1, show_frame)

def faceDetect(frame):

    # Do face detection
    #faces = faceCascade.detectMultiScale(frame, 1.1, 3, 0, (10, 10))

    #Slower method
    gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
    gray = cv2.equalizeHist( gray )
    faces = faceCascade.detectMultiScale(
	    gray,
        scaleFactor=1.1,
        minNeighbors=4,
        minSize=(20, 20),
        flags=cv2.cv.CV_HAAR_SCALE_IMAGE | cv2.cv.CV_HAAR_FIND_BIGGEST_OBJECT | cv2.cv.CV_HAAR_DO_ROUGH_SEARCH
    )

    #print "Found {0} faces!".format(len(faces))
    global last_image_faces
    image_faces = []

    for (x, y, w, h) in faces:
        center_x = x + (w/2)
        center_y = y + (h/2)
        start_y = center_y - 40
        start_x = center_x - 40
        face = gray[y:(y+h), x:(x+w)]
        image_faces.append(face)
        # Draw a green rectangle around the face
        cv2.rectangle(frame, (x, y), (start_x+w, start_y+h), (0, 255, 0), 2)
    last_image_faces = image_faces
    return frame

def saveFace():
    global last_image_faces
    for face in last_image_faces:
        face = cv2.resize(face, (120, 120))
        face = cv2.equalizeHist(face)
        cv2.imwrite("/home/pi/photos/faces/face-" + str(len(face_db) + 1) + ".jpg", face)
        print("Saved image " + str(len(face_db)))
    loadFaceDB()

def buttonPress():
    global saved
    if(GPIO.input(12) == 0):
        if saved == False:
            saveFace()
            saved = True

    else:
        saved = False
    if(GPIO.input(16) == 0):
        if saved == False:
            saveFace()
            saved = True
    else:
        saved = False
    if(GPIO.input(18) == 0):
        if saved == False:
            saveFace()
            saved = True
    else:
        saved = False

def loadFaceDB():
    # Load faces
    face_db_path='/home/pi/photos/faces'
    onlyfiles = [ f for f in listdir(face_db_path) if isfile(join(face_db_path,f)) ]
    global face_db
    face_db = numpy.empty(len(onlyfiles), dtype=object)
    for n in range(0, len(onlyfiles)):
      face_db[n] = cv2.imread( join(face_db_path,onlyfiles[n]) )

loadFaceDB()
show_frame()
root.mainloop()
GPIO.cleanup()