import gradio as gr
import face_alignment
from PIL import Image
import numpy as np
import cv2
import tempfile

# Load model
fa = face_alignment.FaceAlignment('2D', device='cpu', flip_input=False)

def process_mask(mask_img):
    """Automatically remove white/transparent background from mask image"""
    mask_img = mask_img.convert("RGBA")
    datas = mask_img.getdata()
    new_data = []
    white_threshold = 240
    
    for item in datas:
        if (item[0] > white_threshold and 
            item[1] > white_threshold and 
            item[2] > white_threshold) or item[3] == 0:
            new_data.append((255, 255, 255, 0))
        else:
            new_data.append(item)
    
    mask_img.putdata(new_data)
    return mask_img

# Load and process mask
original_mask = Image.open("mask.jpg").convert("RGBA")
clean_mask = process_mask(original_mask)

def fit_mask_to_face(image):
    """Apply mask to the face in the image"""
    img_np = np.array(image)
    landmarks = fa.get_landmarks(img_np)

    if not landmarks:
        return image

    image_rgb = cv2.cvtColor(img_np, cv2.COLOR_BGR2RGB)

    for landmark in landmarks:
        left_eye = landmark[36:42]
        right_eye = landmark[42:48]

        left_pupil = np.mean(left_eye, axis=0).astype(int)
        right_pupil = np.mean(right_eye, axis=0).astype(int)

        eye_distance = np.linalg.norm(right_pupil - left_pupil).astype(int)
        scale_w = int(eye_distance * 3)
        scale_h = int(scale_w * clean_mask.size[1] / clean_mask.size[0])
        resized_mask = clean_mask.resize((scale_w, scale_h), Image.Resampling.LANCZOS)

        center_x = int((left_pupil[0] + right_pupil[0]) / 2)
        center_y = int((left_pupil[1] + right_pupil[1]) / 2) - 20

        top_left_x = center_x - scale_w // 2
        top_left_y = center_y - scale_h // 2

        background_pil = Image.fromarray(image_rgb).convert("RGBA")
        background_pil.paste(resized_mask, (top_left_x, top_left_y), resized_mask)

        return background_pil.convert("RGB")
    return image

def capture_and_process():
    """Capture image from webcam and process it"""
    temp_file = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
    temp_path = temp_file.name
    
    cap = cv2.VideoCapture(0)
    ret, frame = cap.read()
    
    if ret:
        cv2.imwrite(temp_path, frame)
        pil_image = Image.open(temp_path)
        result = fit_mask_to_face(pil_image)
        cap.release()
        return result
    else:
        cap.release()
        return None

# Create interface without CSS
with gr.Blocks() as demo:
    gr.Markdown("# 🎭 Party Face Mask App")
    gr.Markdown("Upload an image or capture from your webcam to try on our party mask!")
    
    with gr.Tabs():
        with gr.TabItem("📁 Upload Image"):
            with gr.Row():
                with gr.Column():
                    upload_input = gr.Image(label="Upload your photo", type="pil")
                    upload_button = gr.Button("Apply Mask")
                with gr.Column():
                    upload_output = gr.Image(label="Result", type="pil")
            upload_button.click(fn=fit_mask_to_face, inputs=upload_input, outputs=upload_output)
            
        with gr.TabItem("📸 Webcam"):
            with gr.Row():
                with gr.Column():
                    webcam_button = gr.Button("Capture Photo from Webcam")
                with gr.Column():
                    webcam_output = gr.Image(label="Result", type="pil")
            webcam_button.click(fn=capture_and_process, outputs=webcam_output)

demo.launch()