import streamlit as st
import face_recognition
import os
import cv2
import insightface
import pickle

from insightface.app import FaceAnalysis


PATH_TMP = 'tmp'
PATH_MODEL = 'bin'
DATA_IMAGE_PICKLE = 'data_images.pkl'
ONNX_SWAPPER_MODEL = 'inswapper_128.onnx'


def face_swapper(image_background, image_customer):
    app = FaceAnalysis(name='buffalo_l')
    app.prepare(ctx_id=0, det_size=(640, 640))
    swapper = insightface.model_zoo.get_model(os.path.join(os.getcwd(), PATH_MODEL, ONNX_SWAPPER_MODEL), download=False)
    face_customer = app.get(image_customer)[0]
    faces = app.get(image_background)

    for face in faces:
        image_background = swapper.get(image_background, face, face_customer, paste_back=True)

    return image_background

def process(image):
    with open(os.path.join(os.getcwd(), PATH_MODEL, DATA_IMAGE_PICKLE), 'rb') as file:
        data_images = pickle.load(file)

    images_background_encoding, images_background_contents = data_images['encodings'], data_images['content']
    image_loaded = face_recognition.load_image_file(image)
    face_encoding = face_recognition.face_encodings(image_loaded)[0]
    face_distances = face_recognition.face_distance(images_background_encoding, face_encoding)

    tmp_distance = face_distances[0]
    tmp_content = images_background_contents[0]
    for face_distance, images_background_content in zip(face_distances[1:], images_background_contents[1:]):
        if tmp_distance > face_distance:
            tmp_distance = face_distance
            tmp_content = images_background_content

    output_image = face_swapper(tmp_content, image_loaded)
    return output_image

image_output = None

st.title('Change Faces')

option = st.radio('How would you like to upload your image?', ('File', 'WebCam'), horizontal=True)

if option=='File':
    uploaded_file = st.file_uploader('Choose your image', type=['jpg', 'png', 'jpeg'])
else:
    uploaded_file = st.camera_input("Take a picture")


if uploaded_file is not None:
    bytes_data = uploaded_file.getvalue()
    if option=='File':
        st.image(uploaded_file)
    if st.button('Process'):
        image_output = process(uploaded_file)
        st.image(image_output)

if image_output is not None:
    image_output_to_download = cv2.cvtColor(image_output, cv2.COLOR_BGR2RGB)
    _, image_output_to_download = cv2.imencode('.jpg', image_output_to_download)
    st.download_button('Download image', image_output_to_download.tobytes(), file_name=f'output_{uploaded_file.name}')