# Dependencies

import os

os.system("pip uninstall -y gradio")
os.system("pip install gradio==3.50")

os.system("pip uninstall -y numpy")
os.system("pip install numpy==1.23.5")


import gradio as gr
import cv2
import insightface
import numpy as np
import gradio.components as gr_comp
import asyncio
import logging
from PIL import Image

# Setup logging
logging.basicConfig(level=logging.INFO)

# Constants for DPI and pixel dimensions
DPI = 300
WIDTH_PIXELS = 1182
HEIGHT_PIXELS = 1773

# Load the insightface model once
providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if cv2.cuda.getCudaEnabledDeviceCount() > 0 else ["CPUExecutionProvider"]
model_path = 'inswapper_128.onnx'
model_swap_insightface = insightface.model_zoo.get_model(model_path, providers=providers)

# Prepare the FaceAnalysis model once
FACE_ANALYSER = insightface.app.FaceAnalysis(
    name="buffalo_l",
    root=".", providers=providers, allowed_modules=["landmark_3d_68", "landmark_2d_106", "detection", "recognition"]
)
FACE_ANALYSER.prepare(ctx_id=0 if cv2.cuda.getCudaEnabledDeviceCount() > 0 else -1, det_size=(640, 640))

# Function to update template choices based on gender
def update_templates(gender):
    if gender == 'Male':
        return gr.Dropdown.update(choices=['Boy Template 1.JPG', 'Boy Template 2.JPG', 'Boy Template 3.JPG'])
    elif gender == 'Female':
        return gr.Dropdown.update(choices=['Girl Template 1.JPG', 'Girl Template 2.JPG', 'Girl Template 3.JPG'])
    else:
        return gr.Dropdown.update(choices=[])

# Main function for face swapping
async def face_swap_and_merge(src_image, gender, template_choice):
    try:
        logging.info("Starting face swap and merge process")
        
        # Resize the source image to the required dimensions for uniform processing
        src_image = cv2.resize(src_image, (WIDTH_PIXELS, HEIGHT_PIXELS))

        template_image_path = f'{template_choice}'
        template_image = cv2.imread(template_image_path, cv2.IMREAD_UNCHANGED)

        # Resize the template image to the required dimensions for uniform processing
        template_image = cv2.resize(template_image, (WIDTH_PIXELS, HEIGHT_PIXELS))

        src_faces = FACE_ANALYSER.get(src_image)
        template_faces = FACE_ANALYSER.get(template_image)

        img_fake = model_swap_insightface.get(img=template_image, target_face=template_faces[0], source_face=src_faces[0], paste_back=True)
        img_fake_rgb = cv2.cvtColor(img_fake, cv2.COLOR_BGR2RGB)

        # Convert to PIL image to set DPI
        pil_image = Image.fromarray(img_fake_rgb)
        pil_image.info['dpi'] = (DPI, DPI)
        
        logging.info("Face swap and merge process completed successfully")
        return np.array(pil_image)
    except Exception as e:
        logging.error(f"Error in face_swap_and_merge: {e}")
        return None

# Function to clear only the output image
def clear_output():
    return None

# Function to clear all inputs and outputs
def clear_all():
    return None, None, None, None

# Create Gradio interface with title and theme
with gr.Blocks(theme='upsatwal/mlsc_tiet') as iface:
    gr.Markdown("# Face Changer")
    gr.Markdown(" ")

    src_image = gr.Image(type="numpy", label="Input Image")
    gender = gr.Dropdown(["Male", "Female"], label="Select Gender", interactive=True)
    template_choice = gr.Dropdown([], label="Select Template", interactive=True)

    gender.change(fn=update_templates, inputs=gender, outputs=template_choice)

    result_image = gr.Image(label="Output Image")

    submit_button = gr.Button("Submit")
    submit_button.click(fn=face_swap_and_merge, inputs=[src_image, gender, template_choice], outputs=result_image)

    clear_output_button = gr.Button("Clear Output")
    clear_output_button.click(fn=clear_output, inputs=[], outputs=result_image)

    clear_all_button = gr.Button("Clear All")
    clear_all_button.click(fn=clear_all, inputs=[], outputs=[src_image, gender, template_choice, result_image])

    iface.launch(debug=True)