Modify working of Image Data Generator to save 50 images for training

Author: msindev

Image_Dataset_Generator.py

@@ -1,9 +1,15 @@
 import cv2
 import numpy as np
 import os
-import time
+import dlib
+from imutils import face_utils
+from imutils.face_utils import FaceAligner
 
-face_cascade = cv2.CascadeClassifier('haarcascades/haarcascade_frontalface_default.xml')
+detector = dlib.get_frontal_face_detector()
+shape_predictor = dlib.shape_predictor("shape_predictor_68_face_landmarks.dat")
+face_aligner = FaceAligner(shape_predictor, desiredFaceWidth=200)
+
+#face_cascade = cv2.CascadeClassifier('haarcascades/haarcascade_frontalface_default.xml')
 
 video_capture = cv2.VideoCapture(0)
 
@@ -17,20 +23,26 @@
 	os.makedirs(directory, exist_ok = 'True')
 
 number_of_images = 0
-MAX_NUMBER_OF_IMAGES = 10
+MAX_NUMBER_OF_IMAGES = 50
 count = 0
 
 while number_of_images < MAX_NUMBER_OF_IMAGES:
 	ret, frame = video_capture.read()
 
 	frame = cv2.flip(frame, 1)
 
-	faces = face_cascade.detectMultiScale(frame, 1.3, 5)
-	for(x,y,w,h) in faces:
-		cv2.rectangle(frame, (x,y),(x+w,y+h),(0,255,0),2)
-		roi = frame[y:y+h, x:x+w]
-		if count == 10:
-			cv2.imwrite(os.path.join(directory, str(name+str(number_of_images)+'.jpg')), roi)
+	frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
+
+	#faces = face_cascade.detectMultiScale(frame, 1.3, 5)
+	faces = detector(frame_gray)
+	if len(faces) == 1:
+		face = faces[0]
+		(x, y, w, h) = face_utils.rect_to_bb(face)
+		face_img = frame_gray[y-50:y + h+100, x-50:x + w+100]
+		face_aligned = face_aligner.align(frame, frame_gray, face)
+
+		if count == 5:
+			cv2.imwrite(os.path.join(directory, str(name+str(number_of_images)+'.jpg')), face_aligned)
 			number_of_images += 1
 			count = 0
 		print(count)