wko_anime-face-similarity/main.py

# Allows imports from the style transfer submodule
import sys
sys.path.append('DCT-Net')

import cv2
import os
import numpy as np

from source.cartoonize import Cartoonizer


def load_source(filename: str) -> np.ndarray:
  return cv2.imread(filename)[...,::-1]


def find_and_crop_face(data: np.ndarray) -> np.ndarray:
  data_gray = cv2.cvtColor(data, cv2.COLOR_BGR2GRAY)
  face_cascade = cv2.CascadeClassifier('haarcascades/haarcascade_frontalface_default.xml')
  face = face_cascade.detectMultiScale(data_gray, 1.3, 4)
  face = max(face, key=len)
  x, y, w, h = face
  face = data[y:y + h, x:x + w]
  return face


def compare_with_anime_characters(data: np.ndarray) -> int:
  # TODO
  return 1


def transfer_to_anime(img: np.ndarray):
  algo = Cartoonizer(dataroot='damo/cv_unet_person-image-cartoon_compound-models')
  return algo.cartoonize(img)


if __name__ == '__main__':
  source = load_source('input.png')
  source_face = find_and_crop_face(source)
  source_face_anime = transfer_to_anime(source)
  print(compare_with_anime_characters(source_face_anime))