import gradio as gr
import cv2
import numpy as np
from PIL import Image
import insightface
from insightface.app import FaceAnalysis

# Initialize face analysis and swapper model
app = FaceAnalysis(name='buffalo_l')
app.prepare(ctx_id=0, det_size=(640, 640))

swapper = insightface.model_zoo.get_model('inswapper_128.onnx')

def to_numpy(img):
    """Convert PIL image to NumPy array in RGB."""
    if isinstance(img, Image.Image):
        return np.array(img.convert("RGB"))
    return img

def enhance_image_colors(image_np):
    """Apply slight contrast and brightness enhancements."""
    lab = cv2.cvtColor(image_np, cv2.COLOR_RGB2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    l_clahe = clahe.apply(l)
    lab_enhanced = cv2.merge((l_clahe, a, b))
    enhanced = cv2.cvtColor(lab_enhanced, cv2.COLOR_LAB2RGB)
    return enhanced

def swap_faces(source_img, target_img, source_index=1, target_index=1):
    # Convert images to NumPy
    source_np = to_numpy(source_img)
    target_np = to_numpy(target_img)

    # Detect faces
    source_faces = sorted(app.get(source_np), key=lambda x: x.bbox[0])
    target_faces = sorted(app.get(target_np), key=lambda x: x.bbox[0])

    if len(source_faces) < source_index or len(target_faces) < target_index:
        raise gr.Error("Face index out of range!")

    source_face = source_faces[source_index - 1]
    target_face = target_faces[target_index - 1]

    # Perform face swap
    result_bgr = swapper.get(target_np, target_face, source_face, paste_back=True)
    result_rgb = cv2.cvtColor(result_bgr, cv2.COLOR_BGR2RGB)

    # Enhance image for better visual quality
    enhanced_result = enhance_image_colors(result_rgb)

    return Image.fromarray(enhanced_result)

# Gradio Interface
gr.Interface(
    fn=swap_faces,
    inputs=[
        gr.Image(label="Source Image"),
        gr.Image(label="Target Image"),
        gr.Number(value=1, label="Source Face Index"),
        gr.Number(value=1, label="Target Face Index")
    ],
    outputs=gr.Image(label="Swapped Output"),
    title="✨ Face Swap with InsightFace",
    description="Upload two images and swap faces using InsightFace (inswapper_128.onnx). Adjust face indices if needed."
).launch()