import gradio as gr
import cv2
import insightface
from insightface.app import FaceAnalysis
import insightface.model_zoo
import matplotlib.pyplot as plt
import requests
import os

file_urls = [
    'https://drive.google.com/file/d/18O9t37O0PUEzIF6yFEIIvoO0ouSyo2a1/view?usp=sharing'
]

def download_file(url, save_name):
    url = url
    if not os.path.exists(save_name):
        file = requests.get(url)
        open(save_name, 'wb').write(file.content)

for i, url in enumerate(file_urls):
    download_file(
        file_urls[i],
        f'inswapper.onnx'
    )

app = FaceAnalysis(name= "buffalo_l")
app.prepare(ctx_id=0, det_size=(640,640))
swapper = insightface.model_zoo.get_model('inswapper.onnx', download=True)

def show_preds_image(image_path1, image_path2):
    poster = cv2.imread(image_path1)
    facial = cv2.imread(image_path2)

    poster_faces = app.get(poster)
    facial_faces = app.get(facial)

    facial_face = facial_faces[0]
    bbox = facial_face['bbox']
    bbox = [int(b) for b in bbox]

    faces = app.get(poster)

    res = poster.copy()
    for face in faces:
        poster = swapper.get(poster, face, facial_face, paste_back=True)

    result_path = 'path_to_save_result_image.jpg'
    cv2.imwrite(result_path, poster)

    return result_path

inputs_image = [
    gr.components.Image(type='filepath', label="Input Image 1"),
    gr.components.Image(type='filepath', label="Input Image 2")
]
output_image = [
    gr.components.Image(type='numpy', label='Output Image')
]
interface_image = gr.Interface(
    fn= show_preds_image,
    inputs= inputs_image,
    outputs= output_image,
    title= "Face Changer",
)

gr.TabbedInterface(
    [interface_image],
    tab_names=['Image interface']
).queue().launch()