import cv2
import numpy as np


def correct_rust_face():
    img_src = cv2.imread('rust_face_src.jpeg')
    pts_src = np.array([[0, 0], [0, 1078], [570, 0], [570, 1078]])

    img_dst = cv2.imread('rust_face_dest.jpeg')
    pts_dst = np.array([[63, 285], [84, 820], [378, 224], [427, 689]])

    dw, dh = img_src.shape[1], img_src.shape[0]

    h, status = cv2.findHomography(pts_dst, pts_src)
    img_out = cv2.warpPerspective(img_dst, h, (dw, dh))
    return img_out


def face_detect(input_image):
    [x, y, w, h] = [0, 0, 0, 0]

    h, w = input_image.shape[:2]
    gray = cv2.cvtColor(input_image, cv2.COLOR_BGR2GRAY)
    gray = cv2.equalizeHist(gray)
    min_size = (w // 10, h // 10)

    face_cascade = cv2.CascadeClassifier(
        "data_haarcascades_haarcascade_frontalface_alt.xml")

    face_rects = face_cascade.detectMultiScale(
        gray, 1.05, 2, cv2.CASCADE_SCALE_IMAGE, min_size)

    if len(face_rects) > 0:
        for face_rect in face_rects:
            x, y, w, h = face_rect
            cv2.rectangle(input_image, (x, y),
                          (x + w, y + h), [255, 255, 0], 2)


if __name__ == '__main__':
    [x, y, w, h] = [0, 0, 0, 0]

    rust_face_src = cv2.imread('rust_face_src.jpeg')
    face_detect(rust_face_src)

    rust_face_dest = cv2.imread('rust_face_dest.jpeg')
    face_detect(rust_face_dest)

    rust_face_dest_correct = correct_rust_face()
    face_detect(rust_face_dest_correct)

    images = np.concatenate(
        (rust_face_src[0: 1070, 0:542, 0:3], rust_face_dest, rust_face_dest_correct[0: 1070, 0:542, 0:3]), axis=1)
    # cv2.imwrite('rust_face_detect.jpeg', images)
    cv2.imshow('rust_face_detect', images)

    cv2.waitKey(0)
    cv2.destroyAllWindows()
