Update app.py

app.py

@@ -20,12 +20,15 @@ DEFAULT_MASK_PATHS = [
     "mask5.png",
 ]
 
-# Verify default masks exist
+# Preload all mask images
+MASK_IMAGES = []
 for mask_path in DEFAULT_MASK_PATHS:
     try:
-        Image.open(mask_path)
+        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; }
@@ -33,20 +36,37 @@ css = """
 #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:
-        # Remove white (RGB > 240), black (RGB < 15), and existing transparent pixels
         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))  # Fully transparent
+            new_data.append((255, 255, 255, 0))
         else:
             new_data.append(item)
     
@@ -67,10 +87,14 @@ def get_random_compliment():
     ]
     return random.choice(compliments)
 
-def apply_mask(person_img, mask_img, seed, randomize_seed):
-    """Improved mask application using facial landmarks with perfect transparency"""
-    if person_img is None or mask_img is None:
-        return None, None, "Empty image"
+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)
@@ -79,8 +103,10 @@ def apply_mask(person_img, mask_img, seed, randomize_seed):
         # Convert to numpy array
         img_np = np.array(person_img)
         
-        # Process the mask with aggressive background removal
+        # 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)
@@ -125,23 +151,24 @@ def apply_mask(person_img, mask_img, seed, randomize_seed):
         print(f"Error applying mask: {e}")
         return person_img, seed, f"Error: {e}"
 
-def capture_and_process(mask_img, seed, randomize_seed):
-    """Capture from webcam and apply mask with transparency"""
-    temp_file = tempfile.NamedTemporaryFile(suffix=".jpg", delete=False)
-    temp_path = temp_file.name
+def capture_and_process(mask_index, seed, randomize_seed):
+    """Capture from webcam and apply selected mask"""
+    if mask_index is None or mask_index >= len(MASK_IMAGES):
+        return None, None, "Invalid mask selection"
+
+    mask_img = MASK_IMAGES[mask_index]
+    if mask_img is None:
+        return None, None, "Mask not found"
 
     cap = cv2.VideoCapture(0)
     ret, frame = cap.read()
-
-    if ret:
-        cv2.imwrite(temp_path, frame)
-        pil_image = Image.open(temp_path)
-        result_img, final_seed, info = apply_mask(pil_image, mask_img, seed, randomize_seed)
-        cap.release()
-        return result_img, final_seed, info
-    else:
-        cap.release()
+    cap.release()
+    
+    if not ret:
         return None, seed, "Failed to capture image"
+    
+    pil_image = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
+    return apply_mask(pil_image, mask_index, seed, randomize_seed)
 
 # Create interface
 with gr.Blocks(css=css) as MaskSuggest:
@@ -158,13 +185,23 @@ with gr.Blocks(css=css) as MaskSuggest:
             )
             
         with gr.Column(elem_id="col-mid"):
-            gr.Markdown("### Step 2. Upload a mask image")
-            mask_img_input = gr.Image(label="Mask image (PNG with transparency)", type="pil")
-            mask_img_example = gr.Examples(
-                examples=["mask1.jpg", "mask2.png", "mask3.png"],
-                inputs=mask_img_input,
-                examples_per_page=12
-            )
+            gr.Markdown("### Step 2. Choose a mask")
+            
+            # Create mask selection thumbnails
+            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)
+                        gr.Image(value=thumbnail, label=f"Mask {i+1}", 
+                                interactive=False, show_label=True,
+                                elem_classes="mask-thumbnail").select(
+                                    lambda x, i=i: i, 
+                                    outputs=[gr.Radio(choices=list(range(len(MASK_IMAGES))), 
+                                              value=0, visible=False)]
+            
+            # Hidden radio buttons for mask selection
+            mask_selection = gr.Radio(choices=list(range(len(MASK_IMAGES))), 
+                                    label="Selected Mask", visible=False, value=0)
             
         with gr.Column(elem_id="col-right"):
             gr.Markdown("### Step 3. Get result")
@@ -173,13 +210,25 @@ with gr.Blocks(css=css) as MaskSuggest:
             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("Run", elem_id="button")
+            process_button = gr.Button("Apply Mask", elem_id="button")
 
     process_button.click(
         fn=apply_mask,
-        inputs=[person_img_input, mask_img_input, seed, randomize_seed],
+        inputs=[person_img_input, mask_selection, seed, randomize_seed],
         outputs=[image_out, seed_used, result_info],
         api_name='apply_mask'
     )
+    
+    gr.Markdown("### Webcam Version")
+    with gr.Row():
+        webcam_button = gr.Button("Capture from Webcam")
+        webcam_output = gr.Image(label="Webcam Result")
+        webcam_status = gr.Textbox(label="Webcam Status")
+    
+    webcam_button.click(
+        fn=capture_and_process,
+        inputs=[mask_selection, seed, randomize_seed],
+        outputs=[webcam_output, seed_used, webcam_status]
+    )
 
 MaskSuggest.queue(api_open=False).launch(show_api=False)