import insightface
import cv2
import gradio as gr
import numpy as np
import os

# Initialize InsightFace model
app = insightface.app.FaceAnalysis(name='buffalo_l', providers=['CPUExecutionProvider'])
app.prepare(ctx_id=0, det_size=(640, 640), det_thresh=0.3)

inswapper_path = "checkpoints/inswapper_128.onnx"
if not os.path.exists(inswapper_path):
    raise FileNotFoundError(f"Model not found at {inswapper_path}")
swapper = insightface.model_zoo.get_model(inswapper_path, providers=['CPUExecutionProvider'])

# Preprocess image to enhance it
def preprocess_image(img):
    alpha = 1.3
    beta = 15
    adjusted = cv2.convertScaleAbs(img, alpha=alpha, beta=beta)
    return adjusted

# Sharpen the image to make it look better
def sharpen_image(img):
    kernel = np.array([[0, -1, 0], [-1, 5, -1], [0, -1, 0]])
    sharpened = cv2.filter2D(img, -1, kernel)
    return np.clip(sharpened, 0, 255).astype(np.uint8)

# Perform face swapping
def swap_faces(src_img, dst_img, blur_strength=5, sharpen=False):
    try:
        # Preprocess images to improve lighting and details
        src = preprocess_image(src_img)
        dst = preprocess_image(dst_img)

        # Convert images to RGB
        src_rgb = cv2.cvtColor(src, cv2.COLOR_BGR2RGB)
        dst_rgb = cv2.cvtColor(dst, cv2.COLOR_BGR2RGB)

        # Detect faces in both images
        src_faces = app.get(src_rgb)
        dst_faces = app.get(dst_rgb)

        if not src_faces or not dst_faces:
            raise ValueError("No faces detected in one or both of the images.")

        # Get the first detected face from each image
        src_face = src_faces[0]
        dst_face = dst_faces[0]

        # Perform the face swap
        swapped_img = swapper.get(dst_rgb, dst_face, src_face, paste_back=True)

        # Apply blur if specified
        if blur_strength > 0:
            swapped_img = cv2.GaussianBlur(swapped_img, (blur_strength, blur_strength), 0)

        # Apply sharpening if specified
        if sharpen:
            swapped_img = sharpen_image(swapped_img)

        # Convert back to BGR for output
        result = cv2.cvtColor(swapped_img, cv2.COLOR_RGB2BGR)
        return result
    except Exception as e:
        print(f"Error: {str(e)}")
        return np.zeros((640, 640, 3), dtype=np.uint8)

# Define Gradio interface
title = "🧠 Futuristic Face Swapper with inswapper_128"
description = (
    "Upload a source face and a target image. The AI swaps the face using inswapper_128.onnx "
    "for clean, smooth results. Adjust blur strength or enable sharpening for enhanced output."
)

demo = gr.Interface(
    fn=swap_faces,
    inputs=[
        gr.Image(label="Source Face", type="numpy"),
        gr.Image(label="Target Image", type="numpy"),
        gr.Slider(label="Blur Strength", minimum=0, maximum=7, step=1, value=5),
        gr.Checkbox(label="Enable Sharpening", value=False)
    ],
    outputs=gr.Image(label="Face Swapped Output"),
    title=title,
    description=description,
    flagging_mode="never"
)

if __name__ == "__main__":
    demo.queue().launch()