factory pattern for face detection

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
+
+        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]
+

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)