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

UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'png', 'jpg', 'jpeg', 'gif', 'mp4', 'avi'}

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(image):
    if image is None:
        return "Please upload an image."
    
    filename = os.path.join(UPLOAD_FOLDER, 'uploaded_image.jpg')
    cv2.imwrite(filename, image)

    faces, img = detect_faces(filename)
    if len(faces) < 2:
        return "Need at least two faces to swap"

    swapped_img = swap_faces(img, faces[0], faces[1])
    
    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 an Image"),
    outputs=gr.Image(type="numpy", label="Swapped Faces"),
    title="Face Swap",
    description="Upload an image with at least two faces, and this tool will swap the faces."
)

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