diff --git a/detect_face.py b/detect_face.py index 54b92bf..d91e006 100644 --- a/detect_face.py +++ b/detect_face.py @@ -1,64 +1,84 @@ +from abc import abstractmethod import cv2 import numpy as np import dlib from matplotlib import pyplot as plt +from abc import ABC +class FaceDetector(ABC): -class FindFaceID: - """ - It takes the image and sends it to the face detection model - by rotating it at 15 degree intervals and returning the original image - according to that angle which has the highest probability of faces in the image. - """ - def __init__(self, detection_method = "ssd", rot_interval = 30) -> None: - self.method = detection_method - self.rot_interval = rot_interval - self.detector = dlib.get_frontal_face_detector() - self.predictor = dlib.shape_predictor("model/shape_predictor_68_face_landmarks.dat") - self.modelFile = "model/res10_300x300_ssd_iter_140000.caffemodel" - self.configFile = "model/deploy.prototxt.txt" - self.FaceNet = cv2.dnn.readNetFromCaffe(self.configFile, self.modelFile) - - def changeOrientationUntilFaceFound(self,image): + @abstractmethod + def changeOrientationUntilFaceFound(self, image, rot_interval): + pass + @abstractmethod + def findFace(self,img): + pass - if(self.method == "dlib"): - rotated_image = self.__searchFaceDlib(image) - return rotated_image - - elif(self.method == "ssd"): - rotated_image = self.__searchFaceSsd(image) - return rotated_image - - else: - print("Select dlib or ssd") - return None + @abstractmethod + def rotate_bound(self,image, angle): + pass - def __findFaceDlib(self, image): - - faces = self.detector(image) - num_of_faces = len(faces) - print("Number of Faces:", num_of_faces ) - if(num_of_faces): - return True - return False + +class DlibFaceDetector(FaceDetector): - def __searchFaceDlib(self, image): + def changeOrientationUntilFaceFound(self, image, rot_interval): img = image.copy() angle_max = 0 - for angle in range(0,360, self.rot_interval): - img_rotated = self.__rotate_bound(img, angle) - is_face_available = self.__findFaceDlib(img_rotated) + for angle in range(0,360, rot_interval): + + img_rotated = self.rotate_bound(img, angle) + is_face_available = self.findFace(img_rotated) + if(is_face_available): return img_rotated return None - - def __searchFaceSsd(self, image): + def findFace(self, image): + + detector = dlib.get_frontal_face_detector() + predictor = dlib.shape_predictor("model/shape_predictor_68_face_landmarks.dat") + faces = detector(image) + num_of_faces = len(faces) + print("Dlib Number of Faces:", num_of_faces ) + if(num_of_faces): + return True + return False + + + def rotate_bound(self, image, angle): + # grab the dimensions of the image and then determine the + # centre + (h, w) = image.shape[:2] + (cX, cY) = (w // 2, h // 2) + + # grab the rotation matrix (applying the negative of the + # angle to rotate clockwise), then grab the sine and cosine + # (i.e., the rotation components of the matrix) + M = cv2.getRotationMatrix2D((cX, cY), angle, 1.0) + cos = np.abs(M[0, 0]) + sin = np.abs(M[0, 1]) + + # compute the new bounding dimensions of the image + nW = int((h * sin) + (w * cos)) + nH = int((h * cos) + (w * sin)) + + # adjust the rotation matrix to take into account translation + M[0, 2] += (nW / 2) - cX + M[1, 2] += (nH / 2) - cY + + # perform the actual rotation and return the image + return cv2.warpAffine(image, M, (nW, nH)) + + +class SsdFaceDetector(FaceDetector): + + + def changeOrientationUntilFaceFound(self,image, rot_interval): """ It takes the image and sends it to the face detection model by rotating it at 15 degree intervals and returning the original image @@ -67,20 +87,40 @@ def __searchFaceSsd(self, image): img = image.copy() face_conf = [] - for angle in range(0,360, self.rot_interval): - img_rotated = self.__rotate_bound(img, angle) - face_conf.append((self.__detectFace(img_rotated), angle)) + for angle in range(0, 360, rot_interval): + img_rotated = self.rotate_bound(img, angle) + face_conf.append((self.findFace(img_rotated), angle)) face_confidence = np.array(face_conf) face_arg_max = np.argmax(face_confidence, axis=0) angle_max = face_confidence[face_arg_max[0]][1] - #print("Maximum face confidence score at angle: ", angle_max) - rotated_img = self.__rotate_bound(image, angle_max) + + rotated_img = self.rotate_bound(image, angle_max) return rotated_img + + def findFace(self,img): + + modelFile = "model/res10_300x300_ssd_iter_140000.caffemodel" + configFile = "model/deploy.prototxt.txt" + FaceNet = cv2.dnn.readNetFromCaffe(configFile, modelFile) - @staticmethod - def __rotate_bound(image, angle): + h, w = img.shape[:2] + blob = cv2.dnn.blobFromImage(cv2.resize(img, (300, 300)), 1.0, + (300, 300), (104.0, 117.0, 123.0)) + + FaceNet.setInput(blob) + faces = FaceNet.forward() + + for i in range(faces.shape[2]): + confidence = faces[0, 0, i, 2] + if confidence > 0.6: + #print("Confidence:", confidence) + return confidence + return 0 + + + def rotate_bound(self,image, angle): # grab the dimensions of the image and then determine the # centre (h, w) = image.shape[:2] @@ -103,18 +143,90 @@ def __rotate_bound(image, angle): # perform the actual rotation and return the image return cv2.warpAffine(image, M, (nW, nH)) + +class HaarFaceDetector(FaceDetector): - def __detectFace(self,img): + def changeOrientationUntilFaceFound(self,image, rot_interval): + img = image.copy() - h, w = img.shape[:2] - blob = cv2.dnn.blobFromImage(cv2.resize(img, (300, 300)), 1.0, - (300, 300), (104.0, 117.0, 123.0)) - self.FaceNet.setInput(blob) - faces = self.FaceNet.forward() + for angle in range(0,360, rot_interval): + + img_rotated = self.rotate_bound(img, angle) + is_face_available = self.findFace(img_rotated) + + if(is_face_available): + return img_rotated - for i in range(faces.shape[2]): - confidence = faces[0, 0, i, 2] - if confidence > 0.6: - #print("Confidence:", confidence) - return confidence - return 0 \ No newline at end of file + + return None + + def findFace(self,img): + + gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) + face_cascade = cv2.CascadeClassifier('model/haarcascade_frontalface_default.xml') + faces = face_cascade.detectMultiScale(gray, 1.3, 5) + num_of_faces = len(faces) + + if(num_of_faces ): + print("Haar Number of Faces:", num_of_faces) + return True + + return False + + + def rotate_bound(self,image, angle): + # grab the dimensions of the image and then determine the + # centre + (h, w) = image.shape[:2] + (cX, cY) = (w // 2, h // 2) + + # grab the rotation matrix (applying the negative of the + # angle to rotate clockwise), then grab the sine and cosine + # (i.e., the rotation components of the matrix) + M = cv2.getRotationMatrix2D((cX, cY), angle, 1.0) + cos = np.abs(M[0, 0]) + sin = np.abs(M[0, 1]) + + # compute the new bounding dimensions of the image + nW = int((h * sin) + (w * cos)) + nH = int((h * cos) + (w * sin)) + + # adjust the rotation matrix to take into account translation + M[0, 2] += (nW / 2) - cX + M[1, 2] += (nH / 2) - cY + + # perform the actual rotation and return the image + return cv2.warpAffine(image, M, (nW, nH)) + +class FaceFactory(ABC): + + @abstractmethod + def get_face_detector(self) -> FaceDetector: + """ Returns new face detector """ + +class DlibModel(FaceFactory): + + def get_face_detector(self) -> FaceDetector: + return DlibFaceDetector() + +class SsdModel(FaceFactory): + + def get_face_detector(self) -> FaceDetector: + return SsdFaceDetector() + +class HaarModel(FaceFactory): + + def get_face_detector(self) -> FaceDetector: + return HaarFaceDetector() + + +def face_factory(face_model = "ssd")->FaceFactory: + """Constructs an face detector factory based on the user's preference.""" + + factories = { + "dlib": DlibModel(), + "ssd" : SsdModel(), + "haar": HaarModel() + } + return factories[face_model] + diff --git a/main.py b/main.py index fc682cd..93764ac 100644 --- a/main.py +++ b/main.py @@ -9,9 +9,9 @@ from extract_words import OcrFactory import extract_words import os -from detect_face import FindFaceID import time import argparse +import detect_face def getCenterRatios(img, centers): """ @@ -134,18 +134,21 @@ def getBoxRegions(regions): parser = argparse.ArgumentParser(description = 'Identity Card Information Extractiion') parser.add_argument('--folder_name', default="images", type=str, help='folder that contain tc id images') parser.add_argument('--neighbor_box_distance', default = 50, type = float, help='Nearest box distance threshold') - parser.add_argument('--face_recognition', default = "ssd", type = str, help='face detection algorithm') - parser.add_argument('--rotation_interval', default = 15, type = int, help='Face search interval for rotation matrix') + parser.add_argument('--face_recognition', default = "haar", type = str, help='face detection algorithm') + parser.add_argument('--ocr_method', default = "EasyOcr", type = str, help='Type of ocr method for converting images to text') + parser.add_argument('--rotation_interval', default = 60, type = int, help='Face search interval for rotation matrix') args = parser.parse_args() Folder = args.folder_name # identity card images folder ORI_THRESH = 3 # Orientation angle threshold for skew correction use_cuda = "cuda" if torch.cuda.is_available() else "cpu" + model = UnetModel("resnet34", use_cuda) nearestBox = NearestBox(distance_thresh = args.neighbor_box_distance, draw_line=False) - findFaceID = FindFaceID(detection_method = args.face_recognition, rot_interval= args.rotation_interval) - Image2Text = extract_words.ocr_factory(ocr_method="EasyOcr", border_thresh=3, denoise = False) + face_detector = detect_face.face_factory(face_model = args.face_recognition) + findFaceID = face_detector.get_face_detector() + Image2Text = extract_words.ocr_factory(ocr_method = args.ocr_method, border_thresh=3, denoise = False) start = time.time() @@ -155,7 +158,7 @@ def getBoxRegions(regions): img = cv2.imread(os.path.join(Folder,filename)) img1 = cv2.cvtColor(img , cv2.COLOR_BGR2RGB) - final_img = findFaceID.changeOrientationUntilFaceFound(img1) + final_img = findFaceID.changeOrientationUntilFaceFound(img1, args.rotation_interval) final_img = utlis.correctPerspective(final_img)