Formatting
This commit is contained in:
parent
a82b98154f
commit
fb84cf1bb5
209
trainer.py
209
trainer.py
@ -7,122 +7,127 @@ import os
|
|||||||
import sys
|
import sys
|
||||||
import json
|
import json
|
||||||
|
|
||||||
|
|
||||||
picCount = 0
|
picCount = 0
|
||||||
new = False
|
new = False
|
||||||
|
|
||||||
|
|
||||||
def start():
|
def start():
|
||||||
print('Starting training..')
|
print('Starting training..')
|
||||||
label = int(input('Label (Integer): '))
|
label = int(input('Label (Integer): '))
|
||||||
name = str(raw_input('Person name: '))
|
name = str(raw_input('Person name: '))
|
||||||
picCount = int(input('Number of Photos: '))
|
picCount = int(input('Number of Photos: '))
|
||||||
global new
|
global new
|
||||||
new = promptNew()
|
new = promptNew()
|
||||||
if new:
|
if new:
|
||||||
newFileName = str(raw_input('New File name: '))
|
newFileName = str(raw_input('New File name: '))
|
||||||
oldFileName = newFileName
|
oldFileName = newFileName
|
||||||
else:
|
else:
|
||||||
print('Found in faces:')
|
print('Found in faces:')
|
||||||
print(os.listdir('faces/'))
|
print(os.listdir('faces/'))
|
||||||
oldFileName = str(raw_input('Old File name: '))
|
oldFileName = str(raw_input('Old File name: '))
|
||||||
newFileName = str(raw_input('New File name: '))
|
newFileName = str(raw_input('New File name: '))
|
||||||
clearTrainFolder()
|
clearTrainFolder()
|
||||||
dum = input('Press any key to start capture..')
|
dum = input('Press any key to start capture..')
|
||||||
capture(picCount)
|
capture(picCount)
|
||||||
addPerson(label, newFileName, oldFileName)
|
addPerson(label, newFileName, oldFileName)
|
||||||
addName(name)
|
addName(name)
|
||||||
|
|
||||||
|
|
||||||
def clearTrainFolder():
|
def clearTrainFolder():
|
||||||
print('clearing train folder')
|
print('clearing train folder')
|
||||||
filelist = [ f for f in os.listdir('train/') if f.endswith(".jpg") ]
|
filelist = [f for f in os.listdir('train/') if f.endswith(".jpg")]
|
||||||
for f in filelist:
|
for f in filelist:
|
||||||
os.remove(os.path.join('train/', f))
|
os.remove(os.path.join('train/', f))
|
||||||
|
|
||||||
|
|
||||||
def addPerson(label, newFileName, oldFileName):
|
def addPerson(label, newFileName, oldFileName):
|
||||||
try:
|
try:
|
||||||
#faceFilePath = '/faces/hesham-saeed2'
|
# faceFilePath = '/faces/hesham-saeed2'
|
||||||
faceFile = 'faces/'+ oldFileName
|
faceFile = 'faces/' + oldFileName
|
||||||
face_recognizer = cv2.createLBPHFaceRecognizer()
|
face_recognizer = cv2.createLBPHFaceRecognizer()
|
||||||
global new
|
global new
|
||||||
if new:
|
if new:
|
||||||
print ('new file. skipping load function')
|
print('new file. skipping load function')
|
||||||
else:
|
else:
|
||||||
print ('loading old file')
|
print('loading old file')
|
||||||
face_recognizer.load(faceFile)
|
face_recognizer.load(faceFile)
|
||||||
except:
|
except:
|
||||||
print('[ Error ] Problem in init/loading LBPHfacerecognizer')
|
print('[ Error ] Problem in init/loading LBPHfacerecognizer')
|
||||||
sys.exit()
|
sys.exit()
|
||||||
images = []
|
images = []
|
||||||
labels = []
|
labels = []
|
||||||
image_names = os.listdir('train/')
|
image_names = os.listdir('train/')
|
||||||
for image_name in image_names:
|
for image_name in image_names:
|
||||||
if image_name.startswith('.'):
|
if image_name.startswith('.'):
|
||||||
continue;
|
continue;
|
||||||
image_path = 'train/' + image_name
|
image_path = 'train/' + image_name
|
||||||
image = cv2.imread(image_path)
|
image = cv2.imread(image_path)
|
||||||
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
|
||||||
face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
|
face_cascade = cv2.CascadeClassifier('haarcascade_frontalface_default.xml')
|
||||||
faces = face_cascade.detectMultiScale(gray, 1.3, 5)
|
faces = face_cascade.detectMultiScale(gray, 1.3, 5)
|
||||||
for (x,y,w,h) in faces:
|
for (x, y, w, h) in faces:
|
||||||
images.append(gray[y:y+h, x:x+w])
|
images.append(gray[y:y + h, x:x + w])
|
||||||
labels.append(label)
|
labels.append(label)
|
||||||
cv2.imshow('face',gray[y:y+h, x:x+w])
|
cv2.imshow('face', gray[y:y + h, x:x + w])
|
||||||
cv2.waitKey(100)
|
cv2.waitKey(100)
|
||||||
print('Learning '+ image_name)
|
print('Learning ' + image_name)
|
||||||
if new:
|
if new:
|
||||||
face_recognizer.train(images, np.array(labels))
|
face_recognizer.train(images, np.array(labels))
|
||||||
else:
|
else:
|
||||||
face_recognizer.update(images, np.array(labels))
|
face_recognizer.update(images, np.array(labels))
|
||||||
face_recognizer.save('faces/'+newFileName)
|
face_recognizer.save('faces/' + newFileName)
|
||||||
print('Updated and saved file in faces/'+newFileName)
|
print('Updated and saved file in faces/' + newFileName)
|
||||||
|
|
||||||
|
|
||||||
def capture(count):
|
def capture(count):
|
||||||
try:
|
try:
|
||||||
print('Initializing camera')
|
print('Initializing camera')
|
||||||
camera = PiCamera()
|
camera = PiCamera()
|
||||||
camera.resolution = (640, 480)
|
camera.resolution = (640, 480)
|
||||||
print('[ OK ] Camera')
|
print('[ OK ] Camera')
|
||||||
sleep(2)
|
sleep(2)
|
||||||
except:
|
except:
|
||||||
print('[ Error ] Can not initialize PiCamera')
|
print('[ Error ] Can not initialize PiCamera')
|
||||||
sys.exit()
|
sys.exit()
|
||||||
for i in range(1, count+1):
|
for i in range(1, count + 1):
|
||||||
pic = 'train/'+str(i)+'.jpg'
|
pic = 'train/' + str(i) + '.jpg'
|
||||||
camera.capture(pic)
|
camera.capture(pic)
|
||||||
photo = cv2.imread(pic,1)
|
photo = cv2.imread(pic, 1)
|
||||||
cv2.imshow('Photo',photo)
|
cv2.imshow('Photo', photo)
|
||||||
print('Captured '+str(i)+'\nPath '+ pic)
|
print('Captured ' + str(i) + '\nPath ' + pic)
|
||||||
cv2.waitKey(100)
|
cv2.waitKey(100)
|
||||||
sleep(1)
|
sleep(1)
|
||||||
cv2.destroyAllWindows()
|
cv2.destroyAllWindows()
|
||||||
print('Done!')
|
print('Done!')
|
||||||
|
|
||||||
|
|
||||||
def promptNew():
|
def promptNew():
|
||||||
filemode = raw_input('Make new File? (y/n): ')
|
filemode = raw_input('Make new File? (y/n): ')
|
||||||
if filemode == 'y':
|
if filemode == 'y':
|
||||||
return True
|
return True
|
||||||
elif filemode == 'n':
|
elif filemode == 'n':
|
||||||
return False
|
return False
|
||||||
else:
|
else:
|
||||||
print('incorrect input')
|
print('incorrect input')
|
||||||
promptNew()
|
promptNew()
|
||||||
|
|
||||||
|
|
||||||
def addName(newName):
|
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()
|
start()
|
||||||
|
|||||||
Loading…
Reference in New Issue
Block a user