-
Notifications
You must be signed in to change notification settings - Fork 3
/
add_person.py
87 lines (73 loc) · 3.18 KB
/
add_person.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
import numpy as np
import cv2
import sys
import os
FREQ_DIV = 5 # frequency divider for capturing training images
RESIZE_FACTOR = 4
NUM_TRAINING = 100
class AddPerson:
def __init__(self):
cascPath = "haarcascades/haarcascade_frontalface_default.xml"
self.face_cascade = cv2.CascadeClassifier(cascPath)
self.face_dir = 'face_data'
self.face_name = sys.argv[1]
self.path = os.path.join(self.face_dir, self.face_name)
if not os.path.isdir(self.face_dir):
os.mkdir(self.face_dir)
if not os.path.isdir(self.path):
os.mkdir(self.path)
self.count_captures = 0
self.count_timer = 0
def capture_training_images(self):
video_capture = cv2.VideoCapture(0)
while True:
self.count_timer += 1
ret, frame = video_capture.read()
inImg = np.array(frame)
outImg = self.process_image(inImg)
cv2.imshow('Video', outImg)
# When everything is done, release the capture on pressing 'q'
if cv2.waitKey(1) & 0xFF == ord('q'):
video_capture.release()
cv2.destroyAllWindows()
return
def process_image(self, inImg):
frame = cv2.flip(inImg, 1)
resized_width, resized_height = (112, 92)
if self.count_captures < NUM_TRAINING:
gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
gray_resized = cv2.resize(gray, (gray.shape[1] / RESIZE_FACTOR, gray.shape[0] / RESIZE_FACTOR))
faces = self.face_cascade.detectMultiScale(
gray_resized,
scaleFactor=1.1,
minNeighbors=5,
minSize=(30, 30),
flags=cv2.CASCADE_SCALE_IMAGE
)
if len(faces) > 0:
areas = []
for (x, y, w, h) in faces:
areas.append(w * h)
max_area, idx = max([(val, idx) for idx, val in enumerate(areas)])
face_sel = faces[idx]
x = face_sel[0] * RESIZE_FACTOR
y = face_sel[1] * RESIZE_FACTOR
w = face_sel[2] * RESIZE_FACTOR
h = face_sel[3] * RESIZE_FACTOR
face = gray[y:y + h, x:x + w]
face_resized = cv2.resize(face, (resized_width, resized_height))
img_no = sorted([int(fn[:fn.find('.')]) for fn in os.listdir(self.path) if fn[0] != '.'] + [0])[-1] + 1
if self.count_timer % FREQ_DIV == 0:
cv2.imwrite('%s/%s.png' % (self.path, img_no), face_resized)
self.count_captures += 1
print "Captured image: ", self.count_captures
cv2.rectangle(frame, (x, y), (x + w, y + h), (0, 255, 0), 3)
cv2.putText(frame, self.face_name, (x - 10, y - 10), cv2.FONT_HERSHEY_PLAIN, 1, (0, 255, 0))
elif self.count_captures == NUM_TRAINING:
print "Training data captured. Press 'q' to exit."
self.count_captures += 1
return frame
if __name__ == '__main__':
trainer = AddPerson()
trainer.capture_training_images()
print "Type in next user to train, or press Recognize"