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)