app.py

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

# Initialize face alignment model
fa = face_alignment.FaceAlignment(face_alignment.LandmarksType.TWO_D, device='cpu', flip_input=False)

# Configuration
MAX_SEED = 999999
DEFAULT_MASK_PATHS = [
    "mask1.jpg",
    "mask2.jpg",
    "mask3.png",
    "mask4.png",
    "mask5.png",
]

# Preload all mask images
MASK_IMAGES = []
for mask_path in DEFAULT_MASK_PATHS:
    try:
        mask_img = Image.open(mask_path)
        MASK_IMAGES.append(mask_img)
    except FileNotFoundError:
        print(f"Warning: Default mask file not found: {mask_path}")
        MASK_IMAGES.append(None)

css = """
#col-left { margin: 0 auto; max-width: 430px; }
#col-mid { margin: 0 auto; max-width: 430px; }
#col-right { margin: 0 auto; max-width: 430px; }
#col-showcase { margin: 0 auto; max-width: 1100px; }
#button { color: blue; }
.mask-thumbnail {
    cursor: pointer;
    border: 2px solid transparent;
    border-radius: 5px;
    margin: 5px;
    height: 100px;
    object-fit: contain;
}
.mask-thumbnail:hover {
    border-color: #4a6cf7;
}
.mask-thumbnail.selected {
    border-color: #4a6cf7;
    box-shadow: 0 0 10px rgba(74, 108, 247, 0.5);
}
"""

def remove_mask_background(mask_img):
    """Remove all backgrounds (white, black, transparent) from mask image"""
    if mask_img is None:
        return None
        
    mask_img = mask_img.convert("RGBA")
    datas = mask_img.getdata()
    new_data = []
    
    for item in datas:
        if (item[0] > 240 and item[1] > 240 and item[2] > 240) or \
           (item[0] < 15 and item[1] < 15 and item[2] < 15) or \
           (item[3] == 0):
            new_data.append((255, 255, 255, 0))
        else:
            new_data.append(item)
    
    mask_img.putdata(new_data)
    return mask_img

def get_random_compliment():
    """Returns a random cute compliment with emoji"""
    compliments = [
        "😍 Wow, you look cute!",
        "✨ You look perfect with this!",
        "🌟 Amazing! This fits you well!",
        "💖 Stunning look!",
        "🎉 Perfect match for your face!",
        "😊 This suits you beautifully!",
        "👌 Flawless execution!",
        "💫 You're rocking this look!"
    ]
    return random.choice(compliments)

def apply_mask(person_img, mask_index, seed, randomize_seed):
    """Apply selected mask to face with perfect transparency"""
    if person_img is None or mask_index is None or mask_index >= len(MASK_IMAGES):
        return None, None, "Empty image or invalid mask"

    mask_img = MASK_IMAGES[mask_index]
    if mask_img is None:
        return None, None, "Mask not found"

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    try:
        # Convert to numpy array
        img_np = np.array(person_img)
        
        # Process the mask with background removal
        processed_mask = remove_mask_background(mask_img)
        if processed_mask is None:
            return None, None, "Mask processing failed"
        
        # Detect facial landmarks
        landmarks = fa.get_landmarks(img_np)

        if not landmarks:
            return person_img, seed, "No face detected"

        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 * processed_mask.size[1] / processed_mask.size[0])
            resized_mask = processed_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

            # Create new RGBA images for perfect transparency handling
            background_pil = Image.fromarray(cv2.cvtColor(img_np, cv2.COLOR_RGB2RGBA))
            resized_mask = resized_mask.convert("RGBA")
            
            # Create transparent result image
            result_img = Image.new("RGBA", background_pil.size)
            result_img.paste(background_pil, (0, 0))
            
            # Paste mask with proper alpha channel
            result_img.paste(resized_mask, (top_left_x, top_left_y), resized_mask)

            return result_img.convert("RGB"), seed, get_random_compliment()
        
        return person_img, seed, "Couldn't apply mask properly"

    except Exception as e:
        print(f"Error applying mask: {e}")
        return person_img, seed, f"Error: {e}"

# Create interface
with gr.Blocks(css=css) as MaskSuggest:
    gr.Markdown("# Mask Suggestion App")
    
    with gr.Row():
        with gr.Column(elem_id="col-left"):
            gr.Markdown("### Step 1. Upload a person image")
            person_img_input = gr.Image(label="Person image", sources='upload', type="pil")
            person_img_example = gr.Examples(
                examples=["human1.jpeg", "human2.jpeg", "human3.jpeg"],
                inputs=person_img_input,
                examples_per_page=12
            )
            
        with gr.Column(elem_id="col-mid"):
            gr.Markdown("### Step 2. Choose a mask")
            
            # Create mask selection thumbnails
            mask_components = []
            with gr.Row():
                for i, mask_img in enumerate(MASK_IMAGES):
                    if mask_img is not None:
                        thumbnail = mask_img.resize((100, 100), Image.Resampling.LANCZOS)
                        thumbnail_component = gr.Image(
                            value=thumbnail,
                            label=f"Mask {i+1}",
                            interactive=False,
                            show_label=True,
                            elem_classes="mask-thumbnail"
                        )
                        mask_components.append(thumbnail_component)
            
            # Hidden radio buttons for mask selection
            mask_selection = gr.Radio(
                choices=list(range(len(MASK_IMAGES))),
                label="Selected Mask",
                visible=False,
                value=0
            )
            
            # Connect thumbnails to radio selection
            for i, component in enumerate(mask_components):
                component.select(
                    fn=lambda x, i=i: i,
                    outputs=[mask_selection]
                )
            
        with gr.Column(elem_id="col-right"):
            gr.Markdown("### Step 3. Get result")
            image_out = gr.Image(label="Result", show_share_button=False)
            seed = gr.Number(value=42, label="Seed", visible=False)
            randomize_seed = gr.Checkbox(value=True, label="Randomize seed", visible=False)
            seed_used = gr.Number(label="Seed used", visible=False)
            result_info = gr.Textbox(label="Status")
            process_button = gr.Button("Apply Mask", elem_id="button")

    process_button.click(
        fn=apply_mask,
        inputs=[person_img_input, mask_selection, seed, randomize_seed],
        outputs=[image_out, seed_used, result_info],
        api_name='apply_mask'
    )

MaskSuggest.queue(api_open=False).launch(show_api=False)