import cv2
import os
import gradio as gr
import numpy as np
from moviepy.editor import *
import logging
os.makedirs('logs',exist_ok=True)
logging.basicConfig(filename='./logs/error.log', level=logging.ERROR)
class FaceSwapper:
    def __init__(self):
        self.logger = logging.getLogger()
        self.face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')

    def swap_faces(self, image1, image2, alpha):
        try:
            # Load the input images
            img1 = cv2.imread(image1.name)
            img2 = cv2.imread(image2.name)
    
            # Detect faces in the images
            faces = self.face_cascade.detectMultiScale(img1)
            faces1 = self.face_cascade.detectMultiScale(img2)
    
            if len(faces) > 0:
                for (x, y, w, h) in faces:
                    # Crop the face region from the first image
                    cropped_face = img1[y:y+h, x:x+w]
    
                    # Calculate dimensions for resizing the face from the second image
                    desired_width = faces1[0][0] + faces1[0][2]
                    desired_height = faces1[0][1] + faces1[0][3]
    
                    # Resize the cropped face to match dimensions
                    resized_cropped_face = cv2.resize(cropped_face, (desired_width, desired_height))
    
                    # Blend the faces using the alpha value
                    blended_face = cv2.addWeighted(resized_cropped_face, alpha, img2[y:y+h, x:x+w], 1-alpha, 0)
    
                    # Position the blended face on the first image
                    x_position = x
                    y_position = y
                    img1[y_position:y_position+desired_height, x_position:x_position+desired_width] = blended_face
    
                return img1
            else:
                return None
                
        except Exception as e:
            logger.error(f"An error occurred during face swapping: {e}")
            return None

    def generate_morph(self, image1, image2, num_frames, alpha):
        try:
            # Load the input images
            img1 = cv2.imread(image1.name)
            img2 = cv2.imread(image2.name)
    
            # Detect faces in the images
            faces = self.face_cascade.detectMultiScale(img1)
            faces1 = self.face_cascade.detectMultiScale(img2)
    
            if len(faces) > 0:
                frames = []
                for i in range(num_frames):
                    alpha_i = alpha + (1-alpha) * i / (num_frames-1)
                    img = self.swap_faces(image1, image2, alpha_i)
                    frames.append(img)
    
                # Convert the frames to a video
                fps = 30
                height, width, _ = frames[0].shape
                fourcc = cv2.VideoWriter_fourcc(*'mp4v')
                video = cv2.VideoWriter('output.mp4', fourcc, fps, (width, height))
    
                for frame in frames:
                    video.write(frame)
    
                video.release()
                return 'output.mp4'
            else:
                return None
                
        except Exception as e:
            logger.error(f"An error occurred during morphing: {e}")
            return None
            
iface = gr.Interface(
    fn=lambda image1, image2, alpha, num_frames: FaceSwapper().generate_morph(image1, image2, num_frames, alpha),
    inputs=[
        gr.inputs.Image(type="filepath", label="Image 1", max_size=1024*1024),  # 1MB limit
        gr.inputs.Image(type="filepath", label="Image 2", max_size=1024*1024),  # 1MB limit
        "slider",
        "number"
    ],
    outputs="video",
    title="Face Swapper",
    description="Swap faces between two images and generate a morph video",
    debug=True,  # enable debug 
    catch_exceptions=True,  # catch exceptions and display error messages instead of restarting the app
)

iface.launch()