Update app.py

app.py

@@ -2,41 +2,135 @@ import gradio as gr
 from rembg import remove
 from PIL import Image
 import numpy as np
+import cv2
+from skimage import filters
+import time
+import os
 
-def remove_background(input_image):
+# Settings
+MAX_SIZE = 1024  # Maximum dimension for input images
+QUALITY = 90     # Output image quality
+
+def enhance_mask(mask):
+    """Refine the mask edges for better quality"""
+    # Convert to grayscale if needed
+    if len(mask.shape) == 3:
+        mask = cv2.cvtColor(mask, cv2.COLOR_BGR2GRAY)
+    
+    # Apply Gaussian blur
+    mask = cv2.GaussianBlur(mask, (5, 5), 0)
+    
+    # Threshold to create binary mask
+    _, binary_mask = cv2.threshold(mask, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)
+    
+    # Apply edge refinement
+    edges = filters.sobel(binary_mask)
+    refined_mask = np.where(edges > 0.1, 255, binary_mask)
+    
+    return refined_mask.astype(np.uint8)
+
+def resize_image(img):
+    """Resize large images while maintaining aspect ratio"""
+    width, height = img.size
+    if max(width, height) > MAX_SIZE:
+        ratio = MAX_SIZE / max(width, height)
+        new_size = (int(width * ratio), int(height * ratio))
+        img = img.resize(new_size, Image.LANCZOS)
+    return img
+
+def remove_background(input_image, post_process=True, alpha_matting=False):
+    start_time = time.time()
+    
     try:
         # Convert to PIL Image if it's a numpy array
         if isinstance(input_image, np.ndarray):
             input_image = Image.fromarray(input_image)
         
+        # Resize if too large
+        input_image = resize_image(input_image)
+        
         # Process with rembg
-        output = remove(input_image)
+        output = remove(
+            input_image, 
+            post_process=post_process,
+            alpha_matting=alpha_matting,
+            alpha_matting_foreground_threshold=240,
+            alpha_matting_background_threshold=10,
+            alpha_matting_erode_size=10
+        )
         
-        # Create mask from alpha channel
+        # Create enhanced mask
         if output.mode == 'RGBA':
             mask = output.split()[-1]
             mask_np = np.array(mask)
+            if post_process:
+                mask_np = enhance_mask(mask_np)
         else:
             mask_np = np.ones(output.size[::-1], dtype=np.uint8) * 255
         
-        return output, Image.fromarray(mask_np)
+        # Apply refined mask
+        if post_process:
+            output.putalpha(Image.fromarray(mask_np))
+        
+        # Calculate processing time
+        proc_time = time.time() - start_time
+        
+        return output, Image.fromarray(mask_np), f"Processed in {proc_time:.2f} seconds"
     
     except Exception as e:
         print(f"Error processing image: {str(e)}")
-        return None, None
+        return None, None, "Error processing image"
+
+# Custom CSS for better UI
+custom_css = """
+.gradio-container { max-width: 900px !important; }
+.output-image { border: 1px solid #e2e8f0 !important; border-radius: 8px !important; }
+.processing-time { font-size: 0.9em; color: #64748b; margin-top: 8px; }
+"""
 
 # Create interface
-iface = gr.Interface(
-    fn=remove_background,
-    inputs=gr.Image(type="pil", label="Input Image"),
-    outputs=[
-        gr.Image(type="pil", label="Result with Transparent Background"), 
-        gr.Image(type="pil", label="Segmentation Mask")
-    ],
-    title="Background Remover (CPU)",
-    description="Upload an image to remove the background using U²-Net. Processing takes 5-15 seconds on CPU."
-)
+with gr.Blocks(css=custom_css) as demo:
+    gr.Markdown("""
+    # 🖼️ Professional Background Remover
+    *Powered by U²-Net with enhanced post-processing*
+    """)
+    
+    with gr.Row():
+        with gr.Column():
+            input_img = gr.Image(label="Upload Image", type="pil", elem_id="input-image")
+            with gr.Accordion("Advanced Options", open=False):
+                post_process = gr.Checkbox(label="Enhanced Post-Processing", value=True)
+                alpha_matting = gr.Checkbox(label="Use Alpha Matting (for fine details)", value=False)
+            submit_btn = gr.Button("Remove Background", variant="primary")
+            
+        with gr.Column():
+            output_img = gr.Image(label="Result", type="pil", elem_id="output-image")
+            output_mask = gr.Image(label="Segmentation Mask", type="pil")
+            time_text = gr.Markdown(elem_classes=["processing-time"])
+    
+    # Examples
+    gr.Examples(
+        examples=[
+            os.path.join(os.path.dirname(__file__), "example1.jpg"),
+            os.path.join(os.path.dirname(__file__), "example2.png")
+        ],
+        inputs=input_img,
+        outputs=[output_img, output_mask, time_text],
+        fn=remove_background,
+        cache_examples=True,
+        label="Try these examples"
+    )
+    
+    submit_btn.click(
+        fn=remove_background,
+        inputs=[input_img, post_process, alpha_matting],
+        outputs=[output_img, output_mask, time_text]
+    )
 
-# Launch with minimal configuration
 if __name__ == "__main__":
-    iface.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=7860,
+        show_error=True,
+        favicon_path=None
+    )