diff --git a/trainer.py b/trainer.py index d72ddbc..5a2e3d5 100644 --- a/trainer.py +++ b/trainer.py @@ -7,122 +7,127 @@ 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) + 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)) + 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) + 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!') + 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() + 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) + 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()