1
+ ### By Bishal Shrestha
2
+
3
+ import cv2
4
+ import os
5
+ import numpy as np
6
+ from PIL import Image
7
+ import pickle
8
+
9
+ BASE_DIR = os .path .dirname (os .path .abspath (__file__ ))
10
+ image_dir = os .path .join (BASE_DIR , "images" )
11
+
12
+ face_cascade = cv2 .CascadeClassifier ('haarcascade_frontalface_default.xml' )
13
+ recognizer = cv2 .face .LBPHFaceRecognizer_create ()
14
+
15
+ current_id = 0
16
+ label_ids = {}
17
+ y_labels = []
18
+ x_train = []
19
+
20
+ for root , dirs , files in os .walk (image_dir ):
21
+ for file in files :
22
+ if file .endswith ("png" ) or file .endswith ("jpg" ):
23
+ path = os .path .join (root , file )
24
+ label = os .path .basename (root ).replace (" " , "-" ).lower ()
25
+ #print(label, path)
26
+ if not label in label_ids :
27
+ label_ids [label ] = current_id
28
+ current_id += 1
29
+ id_ = label_ids [label ]
30
+ #print(label_ids)
31
+ #y_labels.append(label) # some number
32
+ #x_train.append(path) # verify this image, turn into a NUMPY arrray, GRAY
33
+ pil_image = Image .open (path ).convert ("L" ) # grayscale
34
+ size = (550 , 550 )
35
+ final_image = pil_image .resize (size , Image .ANTIALIAS )
36
+ image_array = np .array (final_image , "uint8" )
37
+ #print(image_array)
38
+ faces = face_cascade .detectMultiScale (image_array , scaleFactor = 1.5 , minNeighbors = 5 )
39
+
40
+ for (x ,y ,w ,h ) in faces :
41
+ roi = image_array [y :y + h , x :x + w ]
42
+ x_train .append (roi )
43
+ y_labels .append (id_ )
44
+
45
+
46
+ #print(y_labels)
47
+ #print(x_train)
48
+
49
+ with open ("pickles/face-labels.pickle" , 'wb' ) as f :
50
+ pickle .dump (label_ids , f )
51
+
52
+ recognizer .train (x_train , np .array (y_labels ))
53
+ recognizer .save ("recognizers/face-trainner.yml" )
0 commit comments