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

UPLOAD_FOLDER = 'uploads'

detector = dlib.get_frontal_face_detector()

def detect_faces(image_path):
    img = cv2.imread(image_path)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    faces = detector(gray)
    return faces, img

def extract_face(img, face):
    x, y, w, h = face.left(), face.top(), face.width(), face.height()
    return img[y:y+h, x:x+w]

def swap_faces(img, face1, face2):
    face1_img = extract_face(img, face1)
    face2_img = extract_face(img, face2)

    face1_resized = cv2.resize(face1_img, (face2.width(), face2.height()))
    face2_resized = cv2.resize(face2_img, (face1.width(), face1.height()))

    img[face1.top():face1.top()+face1.height(), face1.left():face1.left()+face1.width()] = face2_resized
    img[face2.top():face2.top()+face2.height(), face2.left():face2.left()+face2.width()] = face1_resized

    return img

def process_file(target_image, source_image):
    if target_image is None or source_image is None:
        return "Please upload both a target image and a source image."
    
    target_filename = os.path.join(UPLOAD_FOLDER, 'target_image.jpg')
    source_filename = os.path.join(UPLOAD_FOLDER, 'source_image.jpg')
    
    cv2.imwrite(target_filename, target_image)
    cv2.imwrite(source_filename, source_image)

    faces_target, img_target = detect_faces(target_filename)
    faces_source, img_source = detect_faces(source_filename)
    
    if len(faces_target) < 1 or len(faces_source) < 1:
        return "Both images must have at least one face."
    
    swapped_img = swap_faces(img_target, faces_target[0], faces_source[0])
    
    result_filename = os.path.join(UPLOAD_FOLDER, 'result_image.jpg')
    cv2.imwrite(result_filename, swapped_img)
    return swapped_img

if not os.path.exists(UPLOAD_FOLDER):
    os.makedirs(UPLOAD_FOLDER)

iface = gr.Interface(
    fn=process_file,
    inputs=[
        gr.Image(type="numpy", label="Upload Target Image"),  # Target image where the face will be swapped
        gr.Image(type="numpy", label="Upload Source Face")  # Source image to extract the face from
    ],
    outputs=gr.Image(type="numpy", label="Swapped Faces"),
    title="Face Swap",
    description="Upload a target image and a source face image, and this tool will swap the face from the source onto the target."
)

if __name__ == "__main__":
    iface.queue().launch(show_error=True, share=True, server_name="0.0.0.0", server_port=7860)