wko_anime-face-similarity/face_detect.py

import cv2
import numpy as np


def equalize_image(data: np.ndarray):
  data_hsv = cv2.cvtColor(data, cv2.COLOR_RGB2HSV)
  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)


def crop_face(data: np.ndarray, bounding_box) -> np.ndarray:
  x, y, w, h = bounding_box
  face = data[y:y + h, x:x + w]
  return face