FaceRecognition-python/trainer.py

from time import sleep
from io import BytesIO
from picamera import PiCamera
import cv2
import numpy as np
import os
import sys
import json


picCount = 0
new = False

def start():
	print('Starting training..')
	label = int(input('Label (Integer): '))
	name = str(raw_input('Person name: '))
	picCount = int(input('Number of Photos: '))
	global new 
	new = promptNew()
	if new:
		newFileName = str(raw_input('New File name: '))
		oldFileName = newFileName
	else:
		print('Found in faces:')
		print(os.listdir('faces/'))
		oldFileName = str(raw_input('Old File name: '))
		newFileName = str(raw_input('New File name: '))
	clearTrainFolder()
	dum = input('Press any key to start capture..')
	capture(picCount)
	addPerson(label, newFileName, oldFileName)
	addName(name)

def clearTrainFolder():
	print('clearing train folder')
	filelist = [ f for f in os.listdir('train/') if f.endswith(".jpg") ]
	for f in filelist:
		os.remove(os.path.join('train/', f))

def addPerson(label, newFileName, oldFileName):
	try:
		#faceFilePath = '/faces/hesham-saeed2'
		faceFile = 'faces/'+ oldFileName
		face_recognizer = cv2.createLBPHFaceRecognizer()
		global new
		if new:
			print 'new file. skipping load function'
		else:
			print 'loading old file'
			face_recognizer.load(faceFile)
	except:
		print('[ Error ] Problem in init/loading LBPHfacerecognizer')
		sys.exit()
	images = []
	labels = []
	image_names = os.listdir('train/')
	for image_name in image_names:
		if image_name.startswith('.'):
			continue;
		image_path = 'train/' + image_name
		image = cv2.imread(image_path)
		gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
		face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
		faces = face_cascade.detectMultiScale(gray, 1.3, 5)
		for (x,y,w,h) in faces:
			images.append(gray[y:y+h, x:x+w])
			labels.append(label)
			cv2.imshow('face',gray[y:y+h, x:x+w])
			cv2.waitKey(100)
			print('Learning '+ image_name)
	if new:
		face_recognizer.train(images, np.array(labels))
	else:
		face_recognizer.update(images, np.array(labels))
	face_recognizer.save('faces/'+newFileName)
	print('Updated and saved file in faces/'+newFileName) 

def capture(count):
	try:
		print('Initializing camera')
		camera = PiCamera()
		camera.resolution = (640, 480)
		print('[ OK ] Camera')
		sleep(2)
	except:
		print('[ Error ] Can not initialize PiCamera')
		sys.exit()
	for i in range(1, count+1):
		pic = 'train/'+str(i)+'.jpg'
		camera.capture(pic)
		photo = cv2.imread(pic,1)
		cv2.imshow('Photo',photo)
		print('Captured '+str(i)+'\nPath '+ pic)
		cv2.waitKey(100)
		sleep(1)
	cv2.destroyAllWindows()
	print('Done!')

def promptNew():
	filemode = raw_input('Make new File? (y/n): ')
	if filemode == 'y':
		return True
	elif filemode == 'n':
		return False
	else:
		print('incorrect input')
		promptNew()

def addName(newName):
	
	with open("faces/names.json", "r") as read_file:
		exist = False
		namesJson = json.load(read_file)
		names = list(namesJson)
		i = 0
		for name in names:
			if newName == name:
				exist = True
				print('name already in names list with index '+str(i))
			i += 1
		if not exist:
			names.append(newName)
			print('Added '+newName+'to names.json file')
	with open("faces/names.json","w") as write_file:
		json.dump(names,write_file)

start()