More robust face detection

face_detect.py

@@ -2,10 +2,29 @@ import cv2
 import numpy as np
 
 
-def find_face_bbox(data: np.ndarray, classifier_file='haarcascades/haarcascade_frontalface_default.xml'):
+def equalize_image(data: np.ndarray):
-  data_gray = cv2.cvtColor(data, cv2.COLOR_RGB2GRAY)
+  data_hsv = cv2.cvtColor(data, cv2.COLOR_RGB2HSV)
-  face_cascade = cv2.CascadeClassifier(classifier_file)
+  data_hsv[:, :, 2] = cv2.equalizeHist(data_hsv[:, :, 2])
+  return cv2.cvtColor(data_hsv, cv2.COLOR_HSV2RGB)
+
+
+def find_face_bbox(data: np.ndarray):
+  classifier_files = [
+    'haarcascades/haarcascade_frontalface_default.xml',
+    'haarcascades/haarcascade_frontalface_alt.xml',
+    'haarcascades/haarcascade_frontalface_alt2.xml',
+    'haarcascades/haarcascade_profileface.xml',
+    'haarcascades/haarcascade_glasses.xml',
+    'lbpcascade_animeface.xml',
+  ]
+  data_equalized = equalize_image(data)
+  data_gray = cv2.cvtColor(data_equalized, cv2.COLOR_RGB2GRAY)
+  face_coords = None
+  for classifier in classifier_files:
+    face_cascade = cv2.CascadeClassifier(classifier)
     face_coords = face_cascade.detectMultiScale(data_gray, 1.1, 3)
+    if face_coords is not None:
+      break
   return max(face_coords, key=len)
 
 

main.py

@@ -1,6 +1,7 @@
 import argparse
 import sys
 import cv2
+import matplotlib.pyplot as plt
 import numpy as np
 
 from metrics import histogram_comparison, structural_similarity_index, euclidean_distance, AccuracyGatherer
@@ -47,13 +48,27 @@ def transfer_to_anime(img: np.ndarray):
   return cv2.cvtColor(model_out, cv2.COLOR_BGR2RGB)
 
 
-def similarity_to_anime(source_image, anime_faces_set):
+def similarity_to_anime(source_image, anime_faces_set, debug=True):
   try:
     source_face_bbox = find_face_bbox(source_image)
   except ValueError:
     return None
   source_anime = transfer_to_anime(source_image)
   source_face_anime = crop_face(source_anime, source_face_bbox)
+
+  if debug:
+    source_image_with_box = source_image.copy()
+    x, y, w, h = source_face_bbox
+    cv2.rectangle(source_image_with_box, (x, y), (x + w, y + h), (255, 0, 0), 2)
+    plt.figure(figsize=[12, 4])
+    plt.subplot(131)
+    plt.imshow(source_image_with_box)
+    plt.subplot(132)
+    plt.imshow(source_anime)
+    plt.subplot(133)
+    plt.imshow(source_face_anime)
+    plt.show()
+
   return compare_with_anime_characters(source_face_anime, anime_faces_set)