wko_anime-face-similarity/main.py
2023-01-29 18:58:45 +01:00

40 lines
1.0 KiB
Python

# 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))