Update app.py

app.py

@@ -8,92 +8,75 @@ print("Initializing LaMa model...")
 lama = SimpleLama(device='cpu')
 
 def ensure_rgb(image):
-    """Convert image to RGB format"""
     if len(image.shape) == 2:
         return cv2.cvtColor(image, cv2.COLOR_GRAY2RGB)
     if len(image.shape) == 3 and image.shape[2] == 4:
         return cv2.cvtColor(image, cv2.COLOR_RGBA2RGB)
     return image
 
-def ensure_mask(mask):
-    """Convert mask to binary grayscale"""
-    if len(mask.shape) == 3:
-        mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
-    _, mask_binary = cv2.threshold(mask, 10, 255, cv2.THRESH_BINARY)
-    return mask_binary
-
+# Function for the UI (Website)
 def process_from_ui(image_dict):
-    """Process image from the UI ImageEditor component"""
     if not image_dict or image_dict["background"] is None:
         return None
-    
     image = ensure_rgb(image_dict["background"])
     
     # Extract mask from layers
     mask = np.zeros(image.shape[:2], dtype=np.uint8)
     if image_dict.get("layers"):
         for layer in image_dict["layers"]:
+            # Gradio 5 layers can be numpy arrays or FileData. 
+            # We assume numpy here as 'type="numpy"' is set in the component.
             if len(layer.shape) == 3:
                 alpha = layer[:, :, 3] if layer.shape[2] == 4 else cv2.cvtColor(layer, cv2.COLOR_RGB2GRAY)
                 mask = cv2.bitwise_or(mask, alpha)
     
-    mask = ensure_mask(mask)
+    _, mask = cv2.threshold(mask, 10, 255, cv2.THRESH_BINARY)
     return lama.predict(image, mask)
 
-def process_api(image, mask):
-    """
-    Simple API endpoint that accepts two separate images.
-    This is what external services (like your Cloudflare Worker) will call.
-    """
+# --- NEW: Function for the API (Cloudflare) ---
+# This accepts two simple images instead of a complex dictionary
+def process_simple_api(image, mask):
     if image is None or mask is None:
         return None
-    
     img_rgb = ensure_rgb(image)
-    mask_binary = ensure_mask(mask)
-    
-    return lama.predict(img_rgb, mask_binary)
+    # Ensure mask is grayscale
+    if len(mask.shape) == 3:
+        mask = cv2.cvtColor(mask, cv2.COLOR_RGB2GRAY)
+    _, mask_bin = cv2.threshold(mask, 10, 255, cv2.THRESH_BINARY)
+    return lama.predict(img_rgb, mask_bin)
 
-# Create the Gradio interface
-with gr.Blocks(title="AI Watermark & Object Remover") as app:
-    gr.Markdown("# 💧 AI Watermark & Object Remover")
-    gr.Markdown("Paint over watermarks or unwanted objects to remove them from your images.")
+with gr.Blocks(title="AI Watermark Remover") as app:
+    gr.Markdown("# 💧 AI Watermark Remover")
     
-    with gr.Tab("Visual Editor"):
-        with gr.Row():
-            input_editor = gr.ImageEditor(
-                label="Draw Mask Over Objects to Remove", 
-                type="numpy",
-                brush=gr.Brush(colors=["#FF0000"], default_size=20),
-                interactive=True
-            )
-            ui_output = gr.Image(label="Result")
+    with gr.Row():
+        # VISUAL UI (For humans)
+        input_editor = gr.ImageEditor(
+            label="Draw Mask", type="numpy", 
+            brush=gr.Brush(colors=["#FF0000"], default_size=20),
+            interactive=True
+        )
+        # HIDDEN INPUTS (For the API / Cloudflare)
+        # We define these so the API is easy to use
+        api_image_input = gr.Image(label="API Image", visible=False, type="numpy")
+        api_mask_input = gr.Image(label="API Mask", visible=False, type="numpy")
         
-        ui_btn = gr.Button("🎨 Remove Objects", variant="primary", size="lg")
-        ui_btn.click(process_from_ui, inputs=input_editor, outputs=ui_output)
+        ui_output = gr.Image(label="Result")
     
-    with gr.Tab("API (For Developers)"):
-        gr.Markdown("""
-        ### API Endpoint
-        Use the `/call/remove_watermark` endpoint with two images:
-        - Upload files to `/upload` first
-        - Then call `/call/remove_watermark` with the file paths
-        
-        Example Cloudflare Worker code is provided separately.
-        """)
-        
-        with gr.Row():
-            api_image = gr.Image(label="Image", type="numpy")
-            api_mask = gr.Image(label="Mask", type="numpy")
-            api_output = gr.Image(label="Result")
-        
-        api_btn = gr.Button("Process (API Test)", variant="secondary")
-        # This creates the API endpoint that can be called via /call/remove_watermark
-        api_btn.click(
-            process_api, 
-            inputs=[api_image, api_mask], 
-            outputs=api_output,
-            api_name="remove_watermark"
-        )
+    ui_btn = gr.Button("Remove Watermark", variant="primary")
+    
+    # UI Button Click
+    ui_btn.click(process_from_ui, inputs=input_editor, outputs=ui_output)
+    
+    # --- THIS IS THE KEY ---
+    # We create a named API endpoint that takes two SIMPLE images.
+    # No "EditorData" dictionary required.
+    api_btn = gr.Button("API_RUN", visible=False)
+    api_btn.click(
+        process_simple_api, 
+        inputs=[api_image_input, api_mask_input], 
+        outputs=ui_output, 
+        api_name="remove_watermark"
+    )
 
 if __name__ == "__main__":
-    app.launch(share=False)
+    app.launch()